#!/usr/bin/env python3
"""Gate G statistical audit: bootstrap 95% confidence intervals (10 000 resamples, seed 0) of median decode throughput per
configuration and of paired/unpaired differences for the pre-registered comparisons. Reads only raw JSONL logs. Writes
benchmarks/processed/audit.json and paper/tables/audit.tex."""
import json, os, glob, random, statistics
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
rng = random.Random(0); B = 10000

def load(path):
    return [json.loads(l) for l in open(path) if l.strip()]

def boot_median(xs):
    meds = []
    for _ in range(B):
        s = [xs[rng.randrange(len(xs))] for _ in xs]; meds.append(statistics.median(s))
    meds.sort(); return statistics.median(xs), meds[int(0.025 * B)], meds[int(0.975 * B)]

def boot_diff(a, b):
    ds = []
    for _ in range(B):
        sa = [a[rng.randrange(len(a))] for _ in a]; sb = [b[rng.randrange(len(b))] for _ in b]
        ds.append(statistics.median(sa) - statistics.median(sb))
    ds.sort(); return statistics.median(a) - statistics.median(b), ds[int(0.025 * B)], ds[int(0.975 * B)]

def sel(recs, **kw):
    out = []
    for r in recs:
        if r.get("warmup"): continue
        if all((r.get(k) == v) if not callable(v) else v(r.get(k)) for k, v in kw.items()): out.append(r)
    return out

comparisons = []
def add(name, a, b, field="gen_tps", unit="tok/s", src=""):
    xa = [r[field] for r in a if r.get(field) is not None]; xb = [r[field] for r in b if r.get(field) is not None]
    if not xa or not xb: return
    ma, la, ha = boot_median(xa); mb, lb, hb = boot_median(xb); d, ld, hd = boot_diff(xa, xb)
    comparisons.append({"comparison": name, "n_a": len(xa), "n_b": len(xb), "median_a": ma, "ci_a": [la, ha], "median_b": mb, "ci_b": [lb, hb],
                        "diff": d, "ci_diff": [ld, hd], "rel_diff_pct": 100 * d / mb if mb else None, "unit": unit, "source": src})

sw = load(f"{ROOT}/benchmarks/raw/gateB/capacity_sweep_laptop/sweep.jsonl")
add("laptop: TITAN all-resident (64 slots, preloaded) vs mlx-lm", sel(sw, backend="titan", capacity=64, tag="warmall"), sel(sw, backend="mlx_lm"), src="gateB/capacity_sweep_laptop/sweep.jsonl")
add("laptop: TITAN 64 slots not preloaded (sync path) vs mlx-lm", sel(sw, backend="titan", capacity=64, tag=""), sel(sw, backend="mlx_lm"), src="gateB/capacity_sweep_laptop/sweep.jsonl")
q = f"{ROOT}/benchmarks/raw/gateC/qwen3_m3u96a/baseline.jsonl"
if os.path.exists(q):
    qq = load(q); add("M3U96a Qwen3-30B: TITAN all-resident (128 slots) vs mlx-lm", sel(qq, backend="titan"), sel(qq, backend="mlx_lm"), src="gateC/qwen3_m3u96a/baseline.jsonl")
cold = load(f"{ROOT}/benchmarks/raw/gateC/cold_nvme_laptop/sweep.jsonl")
for cap in (16, 32):
    base = sel(cold, capacity=cap, predictor="none")
    for pr, m in (("transition", 8), ("transition2", 8), ("transition3", 8), ("transition2", 16), ("prev_rank", 16)):
        add(f"laptop cold NVMe, {cap} slots: {pr} m={m} vs no prefetch", sel(cold, capacity=cap, predictor=pr, m=m), base, src="gateC/cold_nvme_laptop/sweep.jsonl")
attr = load(f"{ROOT}/benchmarks/raw/gateC/cold_nvme_laptop/attribution.jsonl")
base = sel(attr, predictor="none")
for pr, m in (("probe", 8), ("probe", 16), ("probe2", 8), ("probe3", 8), ("transition2", 8), ("transition2t50", 8), ("transition2t70", 8)):
    add(f"laptop cold NVMe, 16 slots (general): {pr} m={m} vs no prefetch", sel(attr, predictor=pr, m=m), base, src="gateC/cold_nvme_laptop/attribution.jsonl")
fz = f"{ROOT}/benchmarks/raw/gateD/fused_upgate_laptop.jsonl"
if os.path.exists(fz):
    f = load(fz)
    add("laptop H4: fused up+gate vs two kernels, 64 slots resident", sel(f, capacity=64, fused_upgate=True), sel(f, capacity=64, fused_upgate=False), src="gateD/fused_upgate_laptop.jsonl")
    add("laptop H4: fused up+gate vs two kernels, 16 slots cold", sel(f, capacity=16, fused_upgate=True), sel(f, capacity=16, fused_upgate=False), src="gateD/fused_upgate_laptop.jsonl")
en = f"{ROOT}/benchmarks/raw/gateG/energy_laptop/bench.jsonl"
if os.path.exists(en):
    e = load(en)
    add("laptop cold NVMe, 16 slots, 256-token generations (energy run): probe m=8 vs no prefetch", sel(e, capacity=16, predictor="probe"), sel(e, capacity=16, predictor="none"), src="gateG/energy_laptop/bench.jsonl")
    add("laptop energy run: TITAN all-resident vs mlx-lm (256 tokens)", sel(e, backend="titan", capacity=64), sel(e, backend="mlx_lm"), src="gateG/energy_laptop/bench.jsonl")
tn = f"{ROOT}/benchmarks/raw/gateE/two_node_m3u96a_m3u96b.jsonl"
if os.path.exists(tn):
    t = load(tn)
    add("M3U96a+M3U96b OLMoE 32 slots: odd experts remote vs local", sel(t, capacity=32, placement="odd"), sel(t, capacity=32, placement="none"), src="gateE/two_node_m3u96a_m3u96b.jsonl")
gf = f"{ROOT}/benchmarks/raw/gateF/two_node_qwen3_m3u96a_m3u96b.jsonl"
if os.path.exists(gf):
    g = load(gf)
    for cap in (32, 64):
        base = sel(g, capacity=cap, placement="none")
        for pl in ("odd", "popular:benchmarks/raw/gateF/traces_qwen3_m3u96a"):
            add(f"M3U96a+M3U96b Qwen3-30B {cap} slots cold: {pl.split(':')[0]} placement vs local cold NVMe", sel(g, capacity=cap, placement=pl), base, src="gateF/two_node_qwen3_m3u96a_m3u96b.jsonl")
os.makedirs(f"{ROOT}/benchmarks/processed", exist_ok=True)
json.dump({"B": B, "seed": 0, "comparisons": comparisons}, open(f"{ROOT}/benchmarks/processed/audit.json", "w"), indent=1)
os.makedirs(f"{ROOT}/paper/tables", exist_ok=True)
L = ["\\begin{table}[t]", "\\centering\\scriptsize", "\\caption{Statistical audit: medians with bootstrap 95\\,\\% CIs (10\\,000 resamples) of decode throughput (tok/s) and of the difference A $-$ B.}", "\\label{tab:audit}",
     "\\begin{tabular}{L{7.2cm}rrrr}", "\\toprule", "Comparison (A vs B) & $n_A/n_B$ & A [CI] & B [CI] & A$-$B [CI] \\\\", "\\midrule"]
for c in comparisons:
    L.append(f"{c['comparison'].replace('_', '\\_').replace('%', '\\%')} & {c['n_a']}/{c['n_b']} & {c['median_a']:.1f} [{c['ci_a'][0]:.1f}, {c['ci_a'][1]:.1f}] & {c['median_b']:.1f} [{c['ci_b'][0]:.1f}, {c['ci_b'][1]:.1f}] & {c['diff']:+.1f} [{c['ci_diff'][0]:+.1f}, {c['ci_diff'][1]:+.1f}] \\\\")
L += ["\\bottomrule", "\\end{tabular}", "\\end{table}"]
open(f"{ROOT}/paper/tables/audit.tex", "w").write("\n".join(L) + "\n")
for c in comparisons:
    print(f"{c['comparison'][:70]:70s} A={c['median_a']:6.1f} [{c['ci_a'][0]:6.1f},{c['ci_a'][1]:6.1f}]  B={c['median_b']:6.1f}  diff={c['diff']:+6.1f} [{c['ci_diff'][0]:+6.1f},{c['ci_diff'][1]:+6.1f}]")