"""
REPRODUCE ALL PAPER RESULTS
============================
Runs every validated experiment from the paper and checks results
against the reported figures. Expected runtime: ~15 minutes.

Usage:
    python3 reproduce_results.py

All results are printed to stdout and also saved to results_check.json.
Global seed: 2025.  Per-trial seeds derived from trial index.
"""
import numpy as np
import json, time
from scipy.stats import ttest_rel

from simulation_utils import (
    make_packet, make_ltf_packet, modulate, demod_bits,
    calc_spectral_efficiency, PILOT_IDX
)
from kr_wifi_receiver import (
    KRReceiver, mmse_equalise, distortion_metric,
    soft_gate, fit_A_sat_gss, pa_rapp_inverse
)

np.random.seed(2025)

def oamp_net(y, h, sigma_n, n_iter=5, damping=0.85):
    r = mmse_equalise(y, h, sigma_n)
    for _ in range(n_iter):
        r = r + damping * np.conj(h) * (y - h*r) / (np.abs(h)**2 + 1e-10)
    return r

def run_trial(snr_db, pa_sat, seed, receiver):
    y, h, sigma_n, x, bits = make_packet(snr_db, pa_sat, seed)
    r_mmse = mmse_equalise(y, h, sigma_n)
    r_kr, _ = receiver.receive(y, h, sigma_n, x[PILOT_IDX], float(snr_db))
    se_mmse = calc_spectral_efficiency(bits, demod_bits(r_mmse))
    se_kr   = calc_spectral_efficiency(bits, demod_bits(r_kr))
    return se_mmse, se_kr

rx = KRReceiver(estimator='gss')
results = {}
all_pass = True

print("="*65)
print("REPRODUCING PAPER RESULTS (n=1000 per condition for speed)")
print("Paper uses n=3000; expect slight Monte Carlo variation.")
print("="*65)

# ── Figure 2: SNR sweep ──────────────────────────────────────────────
print("\n[Fig 2] SNR sweep, A_sat=0.7")
print(f"  {'SNR':>6} {'Gain (paper)':>14} {'Gain (repro)':>14} {'Check':>8}")
snr_ref = {-5:0.033, 10:0.117, 20:0.262, 40:0.590}
snr_results = {}
for snr in [-5, 10, 20, 40]:
    se_m, se_k = [], []
    for s in range(1000):
        m, k = run_trial(snr, 0.7, s, rx)
        se_m.append(m); se_k.append(k)
    gain = np.mean(se_k) - np.mean(se_m)
    ref = snr_ref[snr]; diff = abs(gain - ref)
    ok = 'OK' if diff < 0.05 else 'WARN'
    if ok != 'OK': all_pass = False
    print(f"  {snr:>6}dB {ref:>14.4f} {gain:>14.4f} {ok:>8}")
    snr_results[snr] = float(gain)
results['snr_sweep'] = snr_results

# ── Clean hardware ───────────────────────────────────────────────────
print("\n[Fig 3] Clean hardware (A_sat=1.5, 8-bit ADC)")
se_m, se_k = [], []
for s in range(1000):
    y, h, sigma_n, x, bits = make_packet(20.0, 1.5, s, adc_bits=8)
    r_m = mmse_equalise(y, h, sigma_n)
    r_k, _ = rx.receive(y, h, sigma_n, x[PILOT_IDX], 20.0)
    se_m.append(calc_spectral_efficiency(bits, demod_bits(r_m)))
    se_k.append(calc_spectral_efficiency(bits, demod_bits(r_k)))
gain_clean = np.mean(se_k) - np.mean(se_m)
_, p_clean = ttest_rel(se_k, se_m)
ok = 'OK' if abs(gain_clean) < 0.02 else 'WARN'
if ok != 'OK': all_pass = False
print(f"  Gain={gain_clean:+.4f} bps/Hz  p={p_clean:.4f}  (paper: ~0, p=0.075)  {ok}")
results['clean_gain'] = float(gain_clean)

# ── PA mismatch ──────────────────────────────────────────────────────
print("\n[Fig 4] PA mismatch: K-R vs OAMP-Net at SNR=20dB")
print(f"  {'A_sat':>8} {'KR gain':>10} {'OAMP gain':>11} {'Winner':>8}")
mismatch = {}
for asat in [0.5, 0.7, 0.9, 1.2, 1.5]:
    se_m, se_k, se_o = [], [], []
    for s in range(800):
        y, h, sigma_n, x, bits = make_packet(20.0, asat, s)
        r_m = mmse_equalise(y, h, sigma_n); r_o = oamp_net(y, h, sigma_n)
        r_k, _ = rx.receive(y, h, sigma_n, x[PILOT_IDX], 20.0)
        se_m.append(calc_spectral_efficiency(bits, demod_bits(r_m)))
        se_k.append(calc_spectral_efficiency(bits, demod_bits(r_k)))
        se_o.append(calc_spectral_efficiency(bits, demod_bits(r_o)))
    gk = np.mean(se_k)-np.mean(se_m); go = np.mean(se_o)-np.mean(se_m)
    w = 'KR' if gk > go+0.01 else ('OAMP' if go > gk+0.01 else 'TIE')
    print(f"  {asat:>8.1f} {gk:>+10.4f} {go:>+11.4f} {w:>8}")
    mismatch[asat] = {'kr': float(gk), 'oamp': float(go)}
results['mismatch'] = mismatch

# ── Modulation order ─────────────────────────────────────────────────
print("\n[Table 4] Modulation order scaling at SNR=30dB, A_sat=0.7")
mod_res = {}
for order in [64, 256, 1024]:
    bps = int(np.log2(order)); adc_b = bps + 2
    se_m, se_k = [], []
    for s in range(800):
        rng = np.random.RandomState(s); sigma_n = 10**(-30/20)
        bits = rng.randint(0, 2, 64*bps); x = modulate(bits, order)
        from simulation_utils import tgax_channel, pa_rapp, adc_quantise
        h = tgax_channel(s); noise = sigma_n*(rng.randn(64)+1j*rng.randn(64))/np.sqrt(2)
        y = adc_quantise(h*pa_rapp(x,0.7)+noise, adc_b)
        r_m = mmse_equalise(y, h, sigma_n)
        rp = r_m[PILOT_IDX]; yp=y[PILOT_IDX]; hp=h[PILOT_IDX]; xp=x[PILOT_IDX]
        D = distortion_metric(yp, hp, xp); al = soft_gate(D)
        if al > 0:
            A_fit = fit_A_sat_gss(rp, xp)
            r_k = r_m + al*(pa_rapp_inverse(r_m, A_fit)-r_m)
        else: r_k = r_m
        se_m.append(calc_spectral_efficiency(bits, demod_bits(r_m, order), order))
        se_k.append(calc_spectral_efficiency(bits, demod_bits(r_k, order), order))
    gain = np.mean(se_k)-np.mean(se_m)
    print(f"  {order}-QAM: gain={gain:+.4f}  (paper: +{[0.640,0.692,0.724][[64,256,1024].index(order)]:.3f})")
    mod_res[order] = float(gain)
results['modulation_order'] = mod_res

# ── Two-stage ────────────────────────────────────────────────────────
print("\n[Table 6] Two-stage deployment, A_sat=0.7, L-LTF scale=0.8")
ts_res = {}
for snr in [20, 25, 30]:
    se_perf, se_ts, se_mmse = [], [], []
    for s in range(600):
        y_he, y_ltf, x_ltf, h, sigma_n, x, bits = make_ltf_packet(snr, 0.7, s, 0.8)
        h_est = y_ltf / x_ltf
        # two-stage
        r_ts, _ = rx.receive(y_he, h_est, sigma_n, x[PILOT_IDX], float(snr))
        # perfect CSI
        r_perf, _ = rx.receive(y_he, h, sigma_n, x[PILOT_IDX], float(snr))
        # MMSE
        r_m = mmse_equalise(y_he, h, sigma_n)
        se_perf.append(calc_spectral_efficiency(bits, demod_bits(r_perf)))
        se_ts.append(calc_spectral_efficiency(bits, demod_bits(r_ts)))
        se_mmse.append(calc_spectral_efficiency(bits, demod_bits(r_m)))
    g_p = np.mean(se_perf)-np.mean(se_mmse); g_ts = np.mean(se_ts)-np.mean(se_mmse)
    rec = 100*g_ts/g_p if g_p>0.001 else 0
    print(f"  SNR {snr}dB: perf={g_p:+.4f} two-stage={g_ts:+.4f} recovery={rec:.1f}%")
    ts_res[snr] = {'perf_gain': float(g_p), 'ts_gain': float(g_ts), 'recovery': float(rec)}
results['two_stage'] = ts_res

with open('/home/claude/final_submission/results_check.json', 'w') as f:
    json.dump(results, f, indent=2)

print("\n" + "="*65)
print(f"OVERALL: {'ALL CHECKS PASSED' if all_pass else 'SOME CHECKS WARNED'}")
print("Results saved to results_check.json")
print("="*65)