#!/usr/bin/env python3
"""
Standalone script to construct a Z=128 QC-LDPC matrix with the DE-optimized
degree distribution [7,4,4,4,4,3,3,3] and run FER validation.

Saves results to data/de_results_z128.json.
"""

import numpy as np
import json
import os
import sys
import time

sys.path.insert(0, os.path.dirname(__file__))

from density_evolution import (
    construct_base_matrix, verify_matrix, validate_matrix,
    build_de_profile, compute_threshold, make_profile,
)
from ldpc_analysis import build_peg_matrix

np.random.seed(42)

Z = 128
DEGREES = [7, 4, 4, 4, 4, 3, 3, 3]
CN_MODE = 'normalized'
ALPHA = 0.875
LAM_B = 0.1

print("=" * 70)
print(f"Z={Z} LDPC Matrix Construction and Validation")
print(f"Degrees: {DEGREES}")
print(f"CN mode: {CN_MODE}, alpha: {ALPHA}")
print("=" * 70)

# Step 1: Construct Z=128 matrix
print("\n--- Step 1: Matrix construction (Z=128, reduced trials) ---")
t0 = time.time()
H_base, girth = construct_base_matrix(DEGREES, z=Z, n_trials=500)
t1 = time.time()
print(f"  Construction time: {t1 - t0:.1f}s")

# Step 2: Verify matrix
print("\n--- Step 2: Matrix verification ---")
t0 = time.time()
checks = verify_matrix(H_base, z=Z)
t1 = time.time()
print(f"  Verification time: {t1 - t0:.1f}s")
print(f"  Z={Z} matrix: girth={girth}, rank={checks['actual_rank']}/{checks['expected_rank']}")
print(f"  Full rank: {checks['full_rank']}, Encodable: {checks['encodable']}")
print(f"  Column degrees: {checks['col_degrees']}")
print(f"  Base matrix:\n{H_base}")

# Step 3: DE threshold (Z-independent, quick check)
print("\n--- Step 3: DE threshold check ---")
profile = build_de_profile(DEGREES, m_base=7)
threshold = compute_threshold(profile, lam_b=LAM_B, z_pop=5000, tol=0.5,
                              cn_mode=CN_MODE, alpha=ALPHA)
print(f"  DE threshold (normalized, alpha={ALPHA}): {threshold:.2f} photons/slot")

# Step 4: FER validation with reduced frames
print("\n--- Step 4: FER validation (n_frames=50) ---")
lam_s_points = [1.5, 2.0, 3.0, 4.0, 5.0, 7.0, 10.0]

# Optimized matrix at Z=128
print("\nRunning optimized Z=128 FER...")
t0 = time.time()
fer_opt = validate_matrix(H_base, lam_s_points, n_frames=50, lam_b=LAM_B,
                          z=Z, cn_mode=CN_MODE, alpha=ALPHA)
t1 = time.time()
print(f"  Optimized FER time: {t1 - t0:.1f}s")

# Also build reference matrices at Z=128 for comparison
print("\nRunning PEG ring Z=128 FER...")
# Construct a PEG-ring style matrix at Z=128
from ldpc_analysis import build_peg_matrix as _build_peg_z32
# The PEG ring base matrix structure is Z-independent; just reconstruct with Z=128 shifts
peg_degrees = [7, 3, 3, 3, 2, 2, 2, 2]
H_peg_128, peg_girth = construct_base_matrix(peg_degrees, z=Z, n_trials=500)
peg_checks = verify_matrix(H_peg_128, z=Z)
print(f"  PEG Z=128: girth={peg_girth}, rank={peg_checks['actual_rank']}/{peg_checks['expected_rank']}")

t0 = time.time()
fer_peg = validate_matrix(H_peg_128, lam_s_points, n_frames=50, lam_b=LAM_B,
                          z=Z, cn_mode=CN_MODE, alpha=ALPHA)
t1 = time.time()
print(f"  PEG FER time: {t1 - t0:.1f}s")

# Print results table
print(f"\n{'lam_s':>8s} {'Opt Z=128':>12s} {'PEG Z=128':>12s}")
print("-" * 36)
for lam_s in lam_s_points:
    f_opt = fer_opt[lam_s]['fer']
    f_peg = fer_peg[lam_s]['fer']
    print(f"{lam_s:8.1f} {f_opt:12.3f} {f_peg:12.3f}")

# Save results
output = {
    'z': Z,
    'cn_mode': CN_MODE,
    'alpha': ALPHA,
    'degrees': DEGREES,
    'de_threshold': float(threshold),
    'matrix': H_base.tolist(),
    'matrix_checks': {
        k: (v.item() if hasattr(v, 'item') else v) for k, v in checks.items()
    },
    'girth': int(girth) if girth < float('inf') else None,
    'fer_results': {
        'lam_s_points': lam_s_points,
        'optimized': {str(k): v for k, v in fer_opt.items()},
        'peg_ring': {str(k): v for k, v in fer_peg.items()},
    },
    'peg_matrix': H_peg_128.tolist(),
    'peg_girth': int(peg_girth) if peg_girth < float('inf') else None,
    'peg_matrix_checks': {
        k: (v.item() if hasattr(v, 'item') else v) for k, v in peg_checks.items()
    },
    'n_frames': 50,
}

out_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), '..', 'data')
os.makedirs(out_dir, exist_ok=True)
out_path = os.path.join(out_dir, 'de_results_z128.json')
with open(out_path, 'w') as f:
    json.dump(output, f, indent=2, default=str)
print(f"\nResults saved to {out_path}")