"""Results-dependent LaTeX tables + macros (called by scripts/make_tables.py). Reads run manifests/metrics under
experiments/runs (local copies), results/raw/*.json and results/analyzed/*.json. Every number is traceable to a run id."""
from __future__ import annotations

import glob, json, math, re
from collections import defaultdict
from pathlib import Path

import numpy as np

OUT = Path("paper/tables")
RUNS = "experiments/runs"


def _runs(prefix):
    return sorted(d for d in glob.glob(f"{RUNS}/{prefix}_seed*") if Path(d, "manifest.json").exists())


def _final_val(run):
    rows = [json.loads(l) for l in open(Path(run) / "metrics.jsonl")]
    last = [r for r in rows if any(k.startswith("val_loss_b") for k in r)]
    return last[-1] if last else {}


def table_seed_matrix(macro):
    """Phase 4 decisive matrix: PRISME vs global recurrence vs dense, 5 seeds, val loss by budget and exact match."""
    systems = [("p4v2_prisme", "PRISME sparse editor"), ("p4v2_gr", "Global recurrence (inspired)"), ("p4v2_dense", "Dense 16-slot writes")]
    rows, raw = [], {}
    for pre, lab in systems:
        runs = [r for r in _runs(pre) if json.load(open(Path(r) / "manifest.json")).get("state") == "completed"]
        if not runs:
            continue
        vals = defaultdict(list); ems = defaultdict(list)
        for r in runs:
            fv = _final_val(r)
            for k, v in fv.items():
                if k.startswith("val_loss_b"):
                    vals[int(k[10:])].append(v)
            fem = json.load(open(Path(r) / "manifest.json")).get("final_exact_match", {})
            for split, d in fem.items():
                for b, e in d.items():
                    ems[(split, int(b))].append(e["accuracy"])
        raw[pre] = {"n_seeds": len(runs), "val": {b: v for b, v in vals.items()}, "em": {f"{s}_b{b}": v for (s, b), v in ems.items()}, "run_ids": [Path(r).name for r in runs]}
        rows.append((lab, len(runs), vals, ems))
    if not rows:
        return
    Path("results/analyzed").mkdir(parents=True, exist_ok=True)
    Path("results/analyzed/seed_matrix_v2.json").write_text(json.dumps(raw, indent=1))
    def ms(v):
        v = np.asarray(v, float)
        return f"{v.mean():.4f} $\\pm$ {v.std(ddof=1):.4f}" if len(v) > 1 else f"{v.mean():.4f}"
    def pct(v):
        v = 100 * np.asarray(v, float)
        return f"{v.mean():.1f} $\\pm$ {v.std(ddof=1):.1f}" if len(v) > 1 else f"{v.mean():.1f}"
    head = r"""\begin{table}[t]\centering\small
\caption{Phase 4 decisive comparison on the frozen 5.6M-parameter synthetic backbone (char-level): held-out
answer-span loss on dev-iid under random forced actions at fixed budgets, and free-generation exact match
(dev-iid, 1,000 stratified items). Mean $\pm$ SD over independent training seeds of the adapter; the backbone,
data, recipe and symmetry-breaking device are identical. Budget 0 is the frozen backbone for every system.}
\label{tab:seedmatrix}
\begin{tabular}{llrrrrr}\toprule
System & seeds & loss b=0 & loss b=1 & loss b=2 & loss b=4 & loss b=8\\\midrule
"""
    body = ""
    for lab, n, vals, ems in rows:
        body += f"{lab} & {n} & " + " & ".join(ms(vals.get(b, [float('nan')])) for b in [0, 1, 2, 4, 8]) + r"\\" + "\n"
    body += r"\midrule" + "\n" + r"\multicolumn{7}{l}{\emph{Exact match (\%), dev-iid}: b=0 / b=1 / b=2 / b=8}\\" + "\n"
    for lab, n, vals, ems in rows:
        cells = [pct(ems.get(("dev_iid", b), [float("nan")])) for b in [0, 1, 2, 8]]
        body += f"{lab} & {n} & \\multicolumn{{5}}{{l}}{{{' / '.join(cells)}}}\\\\\n"
    body += r"\multicolumn{7}{l}{\emph{Exact match (\%), held-out template / depth / compose at b=2}}\\" + "\n"
    for lab, n, vals, ems in rows:
        cells = [pct(ems.get((s, 2), [float("nan")])) for s in ["dev_template", "dev_depth", "dev_compose"]]
        body += f"{lab} & {n} & \\multicolumn{{5}}{{l}}{{{' / '.join(cells)}}}\\\\\n"
    (OUT / "seed_matrix.tex").write_text(head + body + r"\bottomrule\end{tabular}\end{table}" + "\n")
    for pre, lab in systems:
        tag = pre.replace("p4v2_", "vtwo")
        if pre not in raw:   # system not (yet) completed: emit explicit placeholders so the manuscript compiles and the audit flags them
            macro(f"{tag}seeds", "0")
            for bn in ("one", "two", "four", "eight"):
                macro(f"{tag}gainb{bn}", "n/a")
            for e in ("embzero", "embtwo", "embeight"):
                macro(f"{tag}{e}", "n/a")
        if pre in raw:
            v = raw[pre]["val"]
            for b, bn in [(1, "one"), (2, "two"), (4, "four"), (8, "eight")]:
                if 0 in v and b in v:
                    macro(f"{tag}gainb{bn}", f"{100 * (1 - np.mean(v[b]) / np.mean(v[0])):.1f}")
            macro(f"{tag}seeds", str(raw[pre]["n_seeds"]))
            em = raw[pre]["em"]
            if "dev_iid_b0" in em and "dev_iid_b2" in em:
                macro(f"{tag}embzero", f"{100 * np.mean(em['dev_iid_b0']):.1f}"); macro(f"{tag}embtwo", f"{100 * np.mean(em['dev_iid_b2']):.1f}")
                macro(f"{tag}embeight", f"{100 * np.mean(em.get('dev_iid_b8', em['dev_iid_b2'])):.1f}")
    macro("vtwobzeroloss", f"{np.mean(raw[systems[0][0]]['val'][0]):.4f}" if systems[0][0] in raw else "n/a")


def table_routing(macro):
    files = sorted(glob.glob("results/raw/routing_v1ctrl.json")) or sorted(glob.glob("results/raw/routing_*.json"))
    if not files:
        return
    d = json.load(open(files[-1]))
    head = r"""\begin{table}[t]\centering\small
\caption{Routing comparison on dev-iid (tiny scale, editor checkpoint of the locked run): held-out answer-span loss,
executed edits per token and analytic extra operations per token (relative to the backbone). Oracles read the
target and are diagnostic ceilings, not deployable systems.}
\label{tab:routing}
\begin{tabular}{llrrr}\toprule
Policy & setting & loss & edits/token & extra ops\\\midrule
"""
    body = ""
    for r in d["records"]:
        if r["policy"] == "oracle_two_step":
            continue
        setting = ", ".join(f"{k}={v}" for k, v in r.items() if k in ("budget", "threshold", "lam", "kap") and v is not None)
        body += f"{r['policy'].replace('_', ' ')}{' (diag.)' if r.get('diagnostic') else ''} & {setting} & {r['loss']:.4f} & {r['edits_per_token']:.2f} & {r['extra_ops_ratio']:.3f}\\\\\n"
    def pick(pol, **kw):
        for r in d["records"]:
            if r["policy"] == pol and all(r.get(k) == v for k, v in kw.items()):
                return r
        return None
    b0 = pick("budget0")
    if b0:
        macro("routingbzero", f"{b0['loss']:.4f}")
        for pol, tag in [("uniform", "uniform"), ("random_forced", "random"), ("fixed_pair", "pair")]:
            for b in (1, 2, 8):
                r = pick(pol, budget=b)
                if r:
                    macro(f"routing{tag}b{'one' if b == 1 else 'two' if b == 2 else 'eight'}", f"{100 * (r['loss'] / b0['loss'] - 1):+.1f}")
        r = pick("controller_locked")
        if r:
            macro("routingctrllocked", f"{100 * (r['loss'] / b0['loss'] - 1):+.1f}"); macro("routingctrlledits", f"{r['edits_per_token']:.2f}"); macro("routingctrllambda", f"{r['lam']}")
        for b in (1, 8):
            r = pick("oracle1_greedy", budget=b)
            if r:
                macro(f"routingoracleb{'one' if b == 1 else 'eight'}", f"{100 * (r['loss'] / b0['loss'] - 1):+.1f}"); macro(f"routingoracleedits{'one' if b == 1 else 'eight'}", f"{r['edits_per_token']:.2f}")
        best_ent = min((r for r in d["records"] if r["policy"] == "entropy"), key=lambda r: r["loss"], default=None)
        best_norm = min((r for r in d["records"] if r["policy"] == "norm"), key=lambda r: r["loss"], default=None)
        if best_ent: macro("routingentropybest", f"{100 * (best_ent['loss'] / b0['loss'] - 1):+.1f}")
        if best_norm: macro("routingnormbest", f"{100 * (best_norm['loss'] / b0['loss'] - 1):+.1f}")
    two = [r for r in d["records"] if r["policy"] == "oracle_two_step"]
    if two:
        t = two[0]
        body += r"\midrule" + f"\\multicolumn{{5}}{{l}}{{Two-step oracle (diag.): one-step {t['one_step_gain']:.4f}, two-step {t['two_step_gain']:.4f}, greedy-then-best {t['greedy_then_best_gain']:.4f} nats/token}}\\\\\n"
        macro("oracleonestep", f"{t['one_step_gain']:.4f}"); macro("oracletwostep", f"{t['two_step_gain']:.4f}")
    (OUT / "routing.tex").write_text(head + body + r"\bottomrule\end{tabular}\end{table}" + "\n")


def table_controller(macro):
    runs = sorted(glob.glob(f"{RUNS}/controller_p4*_seed*/manifest.json"))
    if not runs:
        return
    head = r"""\begin{table}[t]\centering\small
\caption{Value-controller fit on logged counterfactual labels (tiny scale). Calibration-partition diagnostics:
Huber loss of $\hat g$, sign accuracy, Brier score and ECE of the harm head, one-step oracle opportunity and the
regret of the controller's greedy choice on the exhaustive-audit states (captured fraction $=1-$ regret/opportunity).}
\label{tab:controller}
\begin{tabular}{lrrrrrrr}\toprule
Controller features & labels & Huber & sign acc. & Brier & ECE & opportunity & captured\\\midrule
"""
    body = ""
    for mp in runs:
        m = json.load(open(mp))
        if m.get("state") != "completed":
            continue
        ev = m.get("eval_calib", {})
        feats = "h + entropy" if "hent" in m["run_id"] else "z only"
        body += f"{feats} & {ev.get('n_labels', 0):,} & {ev.get('huber', float('nan')):.4f} & {ev.get('sign_acc', float('nan')):.3f} & {ev.get('harm_brier', float('nan')):.3f} & {ev.get('harm_ece', float('nan')):.3f} & {ev.get('oracle_opportunity_per_state', float('nan')):.4f} & {ev.get('captured_fraction', float('nan')):.3f}\\\\\n"
        macro("ctrl" + ("hent" if "hent" in m["run_id"] else "z") + "signacc", f"{100 * ev.get('sign_acc', float('nan')):.1f}")
    (OUT / "controller.tex").write_text(head + body + r"\bottomrule\end{tabular}\end{table}" + "\n")


def table_labels(macro):
    files = sorted(glob.glob("results/analyzed/labels_labels_v1ctrl*.json")) or sorted(glob.glob("results/analyzed/labels_*.json"))
    if not files:
        return
    d = json.load(open(files[-1]))
    for part in ("train", "calib"):
        if part in d and "audit" in d[part]:
            a = d[part]["audit"]
            macro(f"labels{part}within", f"{100 * a['within_share']:.0f}")
            macro(f"labels{part}allpos", f"{100 * a['frac_states_all_positive']:.0f}")
            macro(f"labels{part}allneg", f"{100 * a['frac_states_all_negative']:.0f}")
            macro(f"labels{part}oracle", f"{a['oracle_one_step_mean']:.4f}")
            macro(f"labels{part}stoponly", f"{a['oracle_state_level_mean']:.4f}")
            macro(f"labels{part}random", f"{a['random_action_mean']:.4f}")
            macro(f"labels{part}nstates", f"{a['n_states']:,}")
            if "random_gain_high_vs_low_base_loss" in a:
                macro(f"labels{part}oraclehigh", f"{a['oracle_gain_high_vs_low_base_loss']['high']:.4f}")
                macro(f"labels{part}randomhigh", f"{a['random_gain_high_vs_low_base_loss']['high']:+.4f}")
                macro(f"labels{part}topdecoracle", f"{a['top_decile_base_loss']['oracle']:.3f}"); macro(f"labels{part}topdecrandom", f"{a['top_decile_base_loss']['random']:+.3f}")
                macro(f"labels{part}fracposhigh", f"{100 * a['frac_actions_positive_high_vs_low']['high']:.0f}"); macro(f"labels{part}fracposlow", f"{100 * a['frac_actions_positive_high_vs_low']['low']:.0f}")


def table_latency(macro):
    files = sorted(glob.glob("results/raw/latency_*.json"))
    if not files:
        return
    d = json.load(open(files[-1]))
    head = r"""\begin{table}[t]\centering\small
\caption{Measured decode latency per generated token (fp32, prompt %d tokens, %d generated, warm-up and
synchronisation), device %s. Median and p95 over generated tokens of three repetitions.}
\label{tab:latency}
\begin{tabular}{llrrrr}\toprule
Policy & budget & batch & median ms & p95 ms & cache bytes\\\midrule
""" % (d["prompt_len"], d["n_new"], d["device"].replace("_", " "))
    body = ""
    for r in d["records"]:
        body += f"{r['policy'].replace('_', ' ')} & {r['budget']} & {r['batch']} & {r['median_ms']:.2f} & {r['p95_ms']:.2f} & {r['cache_bytes']:,}\\\\\n"
    (OUT / "latency.tex").write_text(head + body + r"\bottomrule\end{tabular}\end{table}" + "\n")
    b0 = [r for r in d["records"] if r["policy"] == "budget0" and r["batch"] == 1]
    if b0:
        macro("latencybzero", f"{b0[0]['median_ms']:.2f}")


def make_all(macro):
    table_seed_matrix(macro); table_routing(macro); table_controller(macro); table_labels(macro); table_latency(macro)


def table_routing_seeds(macro):
    """Per-seed routing summary for the v2 PRISME editors: budget 0, best fixed schedule, locked controller, oracle."""
    files = sorted(glob.glob("results/raw/routing_v2s*.json"))
    if not files:
        return
    rows = []
    for f in files:
        d = json.load(open(f)); recs = d["records"]
        def pick(pol, **kw):
            for r in recs:
                if r["policy"] == pol and all(r.get(k) == v for k, v in kw.items()):
                    return r
        b0 = pick("budget0"); fixed = [r for r in recs if r["policy"] in ("random_forced", "uniform", "fixed_pair")]
        best_fixed = min(fixed, key=lambda r: r["loss"]); ctl = pick("controller_locked"); ora = pick("oracle1_greedy", budget=8)
        ent = min((r for r in recs if r["policy"] == "entropy"), key=lambda r: r["loss"]); nrm = min((r for r in recs if r["policy"] == "norm"), key=lambda r: r["loss"])
        seed = Path(f).stem.replace("routing_v2s", "")
        man = sorted(glob.glob(f"{RUNS}/controller_v2s{seed}_seed*/manifest.json"))
        ev = json.load(open(man[-1])).get("eval_calib", {}) if man else {}
        rows.append({"seed": seed, "b0": b0["loss"], "best_fixed": best_fixed["loss"], "best_fixed_name": f"{best_fixed['policy'].replace('_', ' ')} b={best_fixed['budget']}",
                     "entropy": ent["loss"], "norm": nrm["loss"], "ctrl": ctl["loss"], "ctrl_edits": ctl["edits_per_token"], "ctrl_lam": ctl["lam"], "ctrl_kap": ctl["kap"],
                     "oracle8": ora["loss"], "oracle8_edits": ora["edits_per_token"], "sign_acc": ev.get("sign_acc", float("nan")), "captured": ev.get("captured_fraction", float("nan"))})
    Path("results/analyzed/routing_seeds_v2.json").write_text(json.dumps(rows, indent=1))
    head = r"""\begin{table}[t]\centering\small
\caption{Learned routing per adapter seed (v2 backbone, dev-iid, 640 items): held-out answer-span loss of the
frozen backbone, of the best fixed schedule (named), of the best entropy and norm routers, of the locked
controller ($\lambda,\kappa$ selected on calibration rollouts; edits per token in parentheses), and of the
diagnostic greedy one-step oracle at budget 8. Calibration sign accuracy and captured fraction of the oracle
opportunity for each controller.}
\label{tab:routingseeds}
\begin{tabular}{lrlrrrrrr}\toprule
seed & budget 0 & best fixed & entropy & norm & controller (edits) & oracle b=8 (edits) & sign acc. & captured\\\midrule
"""
    body = ""
    for r in rows:
        body += (f"{r['seed']} & {r['b0']:.4f} & {r['best_fixed']:.4f} ({r['best_fixed_name']}) & {r['entropy']:.4f} & {r['norm']:.4f} & "
                 f"{r['ctrl']:.4f} ({r['ctrl_edits']:.2f}) & {r['oracle8']:.4f} ({r['oracle8_edits']:.2f}) & {r['sign_acc']:.3f} & {r['captured']:+.3f}\\\\\n")
    (OUT / "routing_seeds.tex").write_text(head + body + r"\bottomrule\end{tabular}\end{table}" + "\n")
    b0 = np.array([r["b0"] for r in rows]); ctl = np.array([r["ctrl"] for r in rows]); bf = np.array([r["best_fixed"] for r in rows]); ora = np.array([r["oracle8"] for r in rows])
    macro("vtwoctrlseeds", str(len(rows))); macro("vtwoctrlrel", f"{100 * (ctl / b0 - 1).mean():+.1f}"); macro("vtwoctrlrelsd", f"{100 * (ctl / b0 - 1).std(ddof=1) if len(rows) > 1 else 0:.1f}")
    macro("vtwobestfixedrel", f"{100 * (bf / b0 - 1).mean():+.1f}"); macro("vtwooraclerel", f"{100 * (ora / b0 - 1).mean():+.1f}")
    macro("vtwoctrlbetter", str(int((ctl < b0).sum())))


_old_make_all = make_all
def make_all(macro):  # noqa: F811
    _old_make_all(macro); table_routing_seeds(macro)


def table_em_families(macro):
    """Per-family exact match (dev-iid) averaged over seeds, by system and budget, with correct<->incorrect transitions."""
    systems = [("p4v2_prisme", "PRISME"), ("p4v2_gr", "Global rec."), ("p4v2_dense", "Dense")]
    data = {}
    for pre, lab in systems:
        for run in _runs(pre):
            if json.load(open(Path(run) / "manifest.json")).get("state") != "completed":
                continue
            recs = {}
            for b in (0, 1, 2, 8):
                p = Path(run) / f"em_dev_iid_b{b}.json"
                if p.exists():
                    recs[b] = {r["uid"]: r for r in json.load(open(p))}
            if 0 in recs:
                data.setdefault(pre, []).append(recs)
    if not data:
        return
    fams = ["arith", "bind", "fsm", "excl", "sched"]
    head = r"""\begin{table}[t]\centering\small
\caption{Per-family exact match on dev-iid (1,000 stratified items; mean over adapter seeds) at budgets 0, 2 and
8, and item transitions between budget 0 and budget 2 (wrong$\to$right / right$\to$wrong, per 1,000 items).}
\label{tab:emfamilies}
\begin{tabular}{llrrrrrr}\toprule
System & budget & """ + " & ".join(fams) + r""" & transitions ($+$/$-$)\\\midrule
"""
    body = ""
    for pre, lab in systems:
        if pre not in data:
            continue
        for b in (0, 2, 8):
            cells = []
            for fam in fams:
                vals = [np.mean([r["correct"] for r in recs[b].values() if r["family"] == fam]) for recs in data[pre] if b in recs]
                cells.append(f"{100 * np.mean(vals):.1f}" if vals else "n/a")
            tr = ""
            if b == 2:
                plus, minus = [], []
                for recs in data[pre]:
                    if 2 in recs:
                        common = set(recs[0]) & set(recs[2])
                        plus.append(sum(1 for u in common if not recs[0][u]["correct"] and recs[2][u]["correct"]) / len(common) * 1000)
                        minus.append(sum(1 for u in common if recs[0][u]["correct"] and not recs[2][u]["correct"]) / len(common) * 1000)
                if plus:
                    tr = f"{np.mean(plus):.0f} / {np.mean(minus):.0f}"
                    macro(f"{pre.replace('p4v2_', 'vtwo')}transplus", f"{np.mean(plus):.0f}"); macro(f"{pre.replace('p4v2_', 'vtwo')}transminus", f"{np.mean(minus):.0f}")
            body += f"{lab} & {b} & " + " & ".join(cells) + f" & {tr}\\\\\n"
    (OUT / "em_families.tex").write_text(head + body + r"\bottomrule\end{tabular}\end{table}" + "\n")


_old_make_all2 = make_all
def make_all(macro):  # noqa: F811
    _old_make_all2(macro); table_em_families(macro)


def table_pilot(macro):
    """103M pilot: Stage A/B manifests, held-out perplexity by policy (evaluation.run), routing and latency at 103M."""
    import math
    rows = []
    def man(prefix):
        ms = sorted(glob.glob(f"{RUNS}/{prefix}_seed*/manifest.json"))
        return json.load(open(ms[-1])) if ms else None
    pre = man("pretrain103m"); en = man("pretrain_enlarged102m"); ed = man("editor_pilot_103m_lr3e4"); gr = man("gr_pilot_103m")
    if pre:
        macro("pilotbbtokens", f"{pre['train_tokens']/1e6:.1f}"); macro("pilotbbhours", f"{pre['wall_seconds']/3600:.1f}"); macro("pilotbbppl", f"{math.exp(pre['final_val_loss']):.1f}")
        macro("pilotbbpeakgb", f"{pre['peak_memory_bytes']/1e9:.0f}"); macro("pilotbbflops", f"{pre['estimated_train_flops']/1e16:.1f}")
    if en and en.get("state") == "completed":
        macro("pilotenppl", f"{math.exp(en['final_val_loss']):.1f}"); macro("pilotenhours", f"{en['wall_seconds']/3600:.1f}")
    else:
        macro("pilotenppl", "n/a"); macro("pilotenhours", "n/a")
    for m, tag in [(ed, "piloted"), (gr, "pilotgr")]:
        if m:
            fv = _final_val(f"{RUNS}/{m['run_id']}")
            b0 = fv.get("val_loss_b0"); 
            for b in (1, 2, 4, 8):
                macro(f"{tag}gainb{ {1:'one',2:'two',4:'four',8:'eight'}[b] }", f"{100*(1 - fv.get(f'val_loss_b{b}', b0)/b0):.2f}" if b0 else "n/a")
            macro(f"{tag}minutes", f"{m['wall_seconds']/60:.0f}"); macro(f"{tag}peakgb", f"{m['peak_memory_bytes']/1e9:.1f}"); macro(f"{tag}params", f"{m['n_trainable']:,}")
        else:
            for bn in ("one", "two", "four", "eight"):
                macro(f"{tag}gainb{bn}", "n/a")
            macro(f"{tag}minutes", "n/a"); macro(f"{tag}peakgb", "n/a"); macro(f"{tag}params", "n/a")
    # perplexity table from evaluation.run outputs
    evs = sorted(glob.glob("results/raw/eval_*_val.json")) + sorted(glob.glob("results/raw/eval_*103m*_test.json")) + sorted(glob.glob("results/raw/eval_enlarged*_test.json"))
    evs = [f for f in evs if "p4v2" not in f]
    if evs:
        head = r"""\begin{table}[t]\centering\small
\caption{103M pilot: held-out cross-entropy (nats/token) and perplexity on 327,680 FineWeb-Edu validation tokens and
on the LOCKED test split (40 $\times$ 16 $\times$ 512 tokens each) under forced random budgets. Budget 0 of the adapters
equals the backbone bitwise.}
\label{tab:pilotppl}
\begin{tabular}{lllrr}\toprule
System & split & policy & loss & ppl\\\midrule
"""
        body = ""
        for f in evs:
            d = json.load(open(f))
            for pol, r in d.get("perplexity", {}).items():
                if pol == "controller":
                    continue
                body += f"{d['system'].replace('_', ' ')} & {d['split']} & {pol.replace('_', ' ')} & {r['loss']:.4f} & {r['ppl']:.2f}\\\\\n"
                if pol == "forced_b0":
                    macro("pilotppl" + re.sub(r"[^a-z]", "", d["system"]) + d["split"], f"{r['ppl']:.2f}")
                if d["split"] == "test" and pol in ("forced_b2", "forced_b4"):
                    b0 = d["perplexity"]["forced_b0"]["loss"]
                    macro("pilottest" + re.sub(r"[^a-z]", "", d["system"]) + {"forced_b2": "btwo", "forced_b4": "bfour"}[pol], f"{100*(r['loss']/b0-1):+.2f}")
        (OUT / "pilot_ppl.tex").write_text(head + body + r"\bottomrule\end{tabular}\end{table}" + "\n")
    rt = Path("results/raw/routing_103m_val.json")
    if rt.exists():
        d = json.load(open(rt)); recs = d["records"]
        def pick(pol, **kw):
            for r in recs:
                if r["policy"] == pol and all(r.get(k) == v for k, v in kw.items()):
                    return r
        b0 = pick("budget0"); ctl = pick("controller_locked"); ora = pick("oracle1_greedy", budget=8); uni1 = pick("uniform", budget=1)
        if b0 and ctl:
            macro("pilotroutingctrl", f"{100*(ctl['loss']/b0['loss']-1):+.2f}"); macro("pilotroutingctrledits", f"{ctl['edits_per_token']:.2f}")
        if b0 and ora:
            macro("pilotroutingoracle", f"{100*(ora['loss']/b0['loss']-1):+.1f}"); macro("pilotroutingoracleedits", f"{ora['edits_per_token']:.2f}")
        if b0 and uni1:
            macro("pilotroutinguniformone", f"{100*(uni1['loss']/b0['loss']-1):+.2f}")
        head = r"""\begin{table}[t]\centering\small
\caption{103M pilot: routing comparison on validation text (20 $\times$ 8 $\times$ 512 tokens). Oracles are diagnostic.}
\label{tab:pilotrouting}
\begin{tabular}{llrrr}\toprule
Policy & setting & loss & edits/token & extra ops\\\midrule
"""
        body = ""
        for r in recs:
            if r["policy"] == "oracle_two_step":
                continue
            setting = ", ".join(f"{k}={v}" for k, v in r.items() if k in ("budget", "threshold", "lam", "kap") and v is not None)
            body += f"{r['policy'].replace('_', ' ')}{' (diag.)' if r.get('diagnostic') else ''} & {setting} & {r['loss']:.4f} & {r['edits_per_token']:.2f} & {r['extra_ops_ratio']:.3f}\\\\\n"
        (OUT / "pilot_routing.tex").write_text(head + body + r"\bottomrule\end{tabular}\end{table}" + "\n")
    lt = Path("results/raw/latency_103m_mps.json")
    if lt.exists():
        d = json.load(open(lt))
        head = r"""\begin{table}[t]\centering\small
\caption{103M pilot: measured decode latency per generated token on %s (fp32, prompt %d tokens, %d generated, three
repetitions, median and p95), batch 1 and 8.}
\label{tab:pilotlatency}
\begin{tabular}{llrrrr}\toprule
Policy & budget & batch & median ms & p95 ms & cache bytes\\\midrule
""" % (d["device"], d["prompt_len"], d["n_new"])
        body = ""
        for r in d["records"]:
            body += f"{r['policy'].replace('_', ' ')} & {r['budget']} & {r['batch']} & {r['median_ms']:.2f} & {r['p95_ms']:.2f} & {r['cache_bytes']:,}\\\\\n"
        (OUT / "pilot_latency.tex").write_text(head + body + r"\bottomrule\end{tabular}\end{table}" + "\n")
        b0 = [r for r in d["records"] if r["policy"] == "budget0" and r["batch"] == 1]; s8 = [r for r in d["records"] if r["policy"] == "forced_single" and r["budget"] == 8 and r["batch"] == 1]
        d8 = [r for r in d["records"] if r["policy"] == "dense_forced" and r["budget"] == 8 and r["batch"] == 1]; c = [r for r in d["records"] if r["policy"] == "controller" and r["batch"] == 1]
        if b0: macro("pilotlatbzero", f"{b0[0]['median_ms']:.1f}")
        if s8: macro("pilotlatseight", f"{s8[0]['median_ms']:.1f}")
        if d8: macro("pilotlatdeight", f"{d8[0]['median_ms']:.1f}")
        if c: macro("pilotlatctrl", f"{c[0]['median_ms']:.1f}")


_old_make_all3 = make_all
def make_all(macro):  # noqa: F811
    _old_make_all3(macro); table_pilot(macro)


def table_test(macro):
    """Locked test regimes (synthetic): exact match by system and budget, mean ± SD over seeds, plus paired statistics."""
    files = sorted(glob.glob("results/raw/eval_p4v2_*_test_*.json"))
    if not files:
        return
    from evaluation.statistical_analysis import paired_bootstrap, seed_interval
    acc = defaultdict(lambda: defaultdict(list)); items = defaultdict(dict)
    for f in files:
        d = json.load(open(f)); sysname = d["system"]; base = sysname.rsplit("_s", 1)[0]; seed = sysname.rsplit("_s", 1)[1]
        for pol, em in d["exact_match"].items():
            acc[(base, pol)][d["split"]].append(em["accuracy"])
            ip = Path(f).parent / f"items_{sysname}_{d['split']}_{pol}.json"
            if ip.exists():
                items[(base, pol, d["split"])][seed] = {r["uid"]: r["correct"] for r in json.load(open(ip))}
    splits = ["test_iid", "test_template", "test_depth", "test_compose"]
    labels = {"p4v2_prisme": "PRISME sparse editor", "p4v2_gr": "Global recurrence", "p4v2_dense": "Dense 16-slot"}
    head = r"""\begin{table}[t]\centering\small
\caption{LOCKED test regimes (synthetic suite, 1,000 stratified items per regime, unsealed 2026-10-30 after the
pre-registration was frozen): exact match (\%) at budget 0 (frozen backbone, identical for all systems) and at forced
random budget 2, mean $\pm$ SD over adapter seeds. ``pooled'' = the pre-registered primary outcome (iid+template+depth).}
\label{tab:test}
\begin{tabular}{llrrrrr}\toprule
System & budget & iid & template & depth & compose & pooled\\\midrule
"""
    body = ""; pooled = {}
    for base, lab in labels.items():
        for pol in ("budget0", "forced_b2"):
            k = (base, pol)
            if k not in acc:
                continue
            cells = []
            for sp in splits:
                v = np.array(acc[k].get(sp, [np.nan]))
                cells.append(f"{100*v.mean():.1f} $\\pm$ {100*v.std(ddof=1):.1f}" if len(v) > 1 else f"{100*v.mean():.1f}")
            # pooled per seed
            per_seed = []
            n_seeds = min(len(acc[k][sp]) for sp in splits[:3] if sp in acc[k])
            for i in range(n_seeds):
                per_seed.append(np.mean([acc[k][sp][i] for sp in splits[:3]]))
            pooled[k] = per_seed
            pv = np.array(per_seed)
            cells.append(f"{100*pv.mean():.1f} $\\pm$ {100*pv.std(ddof=1):.1f}" if len(pv) > 1 else f"{100*pv.mean():.1f}")
            body += f"{lab} & {0 if pol == 'budget0' else 2} & " + " & ".join(cells) + r"\\" + "\n"
    (OUT / "test.tex").write_text(head + body + r"\bottomrule\end{tabular}\end{table}" + "\n")
    # pre-registered P1: pooled EM differences at b=2 (paired by item within seed, t over seeds)
    stats = {}
    def pooled_items(base, pol, seed):
        out = {}
        for sp in splits[:3]:
            out.update({f"{sp}:{u}": c for u, c in items.get((base, pol, sp), {}).get(seed, {}).items()})
        return out
    for a, b, name in [("p4v2_prisme", None, "prismebtwovsbzero"), ("p4v2_prisme", "p4v2_gr", "prismevsgr"), ("p4v2_prisme", "p4v2_dense", "prismevsdense")]:
        diffs = []
        seeds = sorted({s for (bb, pol, sp), d in items.items() if bb == a and pol == "forced_b2" for s in d})
        for s in seeds:
            x = pooled_items(a, "forced_b2", s)
            y = pooled_items(a, "budget0", s) if b is None else pooled_items(b, "forced_b2", s)
            common = sorted(set(x) & set(y))
            if common:
                diffs.append(paired_bootstrap(np.array([x[u] for u in common]), np.array([y[u] for u in common]))["diff"])
        if diffs:
            si = seed_interval(diffs); stats[name] = si
            macro(f"test{name}", f"{100*si['mean']:+.1f}"); macro(f"test{name}lo", f"{100*si['ci_low']:+.1f}"); macro(f"test{name}hi", f"{100*si['ci_high']:+.1f}"); macro(f"test{name}n", str(si["n"]))
    Path("results/analyzed/test_statistics.json").write_text(json.dumps({"pooled_per_seed": {f"{k[0]}|{k[1]}": v for k, v in pooled.items()}, "stats": stats}, indent=1))


_old_make_all4 = make_all
def make_all(macro):  # noqa: F811
    _old_make_all4(macro); table_test(macro)


def table_controls(macro):
    import torch
    from training.train_editor import build_model
    from training.common import load_yaml
    rows = []
    def count(exp, label, data, compute, regime):
        cfg, m = build_model(exp)
        nb = sum(p.numel() for p in m.backbone.parameters()); tot = sum(p.numel() for p in m.parameters())
        rows.append((label, nb, tot - nb, data, compute, regime))
    with torch.device("meta"):
        count(load_yaml("configs/experiments/p4v2_prisme.yaml"), "Frozen backbone v2 (budget 0)", "synthetic train, 5k steps $\\times$ 64", "backbone only", "no workspace (bitwise bypass)")
        count(load_yaml("configs/experiments/p4v2_prisme.yaml"), "PRISME sparse editor (tiny)", "same + 3k adapter steps $\\times$ 32", "frozen bb fwd + editor", "1--2 slots/action, 0--8 actions")
        count(load_yaml("configs/experiments/p4v2_gr.yaml"), "Global-recurrence adapter (tiny)", "same", "frozen bb fwd + block", "whole latent/action, 0--8 actions")
        count(load_yaml("configs/experiments/p4v2_dense.yaml"), "Dense 16-slot editor (tiny)", "same", "frozen bb fwd + 16 lanes", "16 slots/action, 0--8 actions")
        e = load_yaml("configs/experiments/editor_pilot_103m_lr3e4.yaml"); e["init"] = {}
        count(e, "PRISME-103M sparse editor", "FineWeb-Edu 98.3M tok + 24.6M adapter tok", "$6NP$ + editor", "1--2 slots/action")
        g = load_yaml("configs/experiments/gr_pilot_103m.yaml"); g["init"] = {}
        count(g, "Global-recurrence adapter (103M)", "same", "$6NP$ + block", "whole latent/action")
        count({"model": "configs/baselines/transformer_enlarged_102m.yaml", "mode": "backbone_only"}, "Plain decoder, parameter-matched (102.2M)", "FineWeb-Edu 98.3M tok (same recipe)", "$6NP$", "no workspace")
    with open(OUT / "controls.tex", "w") as f:
        f.write(r"""\begin{table}[t]\centering\small
\caption{Systems and controls. ``bb'' = backbone, ``$6NP$'' = the usual training-FLOP estimate for $N$ tokens and $P$
parameters; adapter training additionally executes up to 8 editor applications per position. Adapter parameters are
unique parameters outside the backbone (counted by instantiation; the 103M editor includes the residual-plus-entropy
controller features).}
\label{tab:controls}
\begin{tabular}{lrrlll}\toprule
System & backbone & adapter & training exposure & training compute & inference regime\\\midrule
""")
        for lab, nb, ad, data, comp, reg in rows:
            f.write(f"{lab} & {nb:,} & {ad:,} & {data} & {comp} & {reg}\\\\\n")
        f.write(r"\bottomrule\end{tabular}\end{table}" + "\n")


_old_make_all5 = make_all
def make_all(macro):  # noqa: F811
    _old_make_all5(macro); table_controls(macro)


def table_ablations(macro):
    """Tiny-scale ablations of the sparse editor (v2 backbone, 3 seeds): slots M, editor depth, pair scheme."""
    systems = [("p4v2_prisme", "Reference: $M=16$, $L_w=2$, adjacent pairs"), ("p4v2_abl_M4", "$M = 4$ slots"), ("p4v2_abl_M8", "$M = 8$ slots"),
               ("p4v2_abl_M32", "$M = 32$ slots"), ("p4v2_abl_L1", "$L_w = 1$ block"), ("p4v2_abl_L4", "$L_w = 4$ blocks"), ("p4v2_abl_pairoffset", "offset pairs (i, i+M/2)")]
    rows = []
    for pre, lab in systems:
        runs = [r for r in _runs(pre) if json.load(open(Path(r) / "manifest.json")).get("state") == "completed"]
        if not runs:
            continue
        vals = defaultdict(list); ems = defaultdict(list); params = []
        for r in runs:
            fv = _final_val(r)
            for k, v in fv.items():
                if k.startswith("val_loss_b"):
                    vals[int(k[10:])].append(v)
            m = json.load(open(Path(r) / "manifest.json")); params.append(m.get("n_trainable", 0))
            for b, e in m.get("final_exact_match", {}).get("dev_iid", {}).items():
                ems[int(b)].append(e["accuracy"])
        rows.append((lab, len(runs), int(np.mean(params)), vals, ems))
    if len(rows) < 2:
        return
    def ms(v, f=4):
        v = np.asarray(v, float); return (f"{v.mean():.{f}f} $\\pm$ {v.std(ddof=1):.{f}f}" if len(v) > 1 else f"{v.mean():.{f}f}")
    def pct(v):
        v = 100 * np.asarray(v, float); return (f"{v.mean():.1f} $\\pm$ {v.std(ddof=1):.1f}" if len(v) > 1 else f"{v.mean():.1f}")
    head = r"""\begin{table}[t]\centering\small
\caption{Ablations of the sparse editor on the frozen v2 synthetic backbone (dev-iid; random forced budgets; mean $\pm$ SD
over adapter seeds; the reference row is the 5-seed system of Table~\ref{tab:seedmatrix}). ``params'' = trainable adapter
parameters excluding the controller.}
\label{tab:ablations}
\begin{tabular}{lrrrrrr}\toprule
Variant & seeds & params & loss b=1 & loss b=2 & loss b=8 & EM b=2 (\%)\\\midrule
"""
    body = ""
    for lab, n, pr, vals, ems in rows:
        body += f"{lab} & {n} & {pr:,} & {ms(vals.get(1, [np.nan]))} & {ms(vals.get(2, [np.nan]))} & {ms(vals.get(8, [np.nan]))} & {pct(ems.get(2, [np.nan]))}\\\\\n"
    (OUT / "ablations.tex").write_text(head + body + r"\bottomrule\end{tabular}\end{table}" + "\n")
    for pre, lab in systems[1:]:
        for r in rows:
            pass
    ref = next((r for r in rows if r[0].startswith("Reference")), None)
    keys = {"Reference": "ref", "$M = 4$": "mfour", "$M = 8$": "meight", "$M = 32$": "mthirtytwo", "$L_w = 1$": "lone", "$L_w = 4$": "lfour", "offset": "offset"}
    for lab, n, pr, vals, ems in rows:
        key = next(v for k, v in keys.items() if lab.startswith(k))
        if 2 in vals and 0 in vals:
            macro("abl" + key + "gainbtwo", f"{100*(1-np.mean(vals[2])/np.mean(vals[0])):.1f}")
        if 8 in vals and 0 in vals:
            macro("abl" + key + "gainbeight", f"{100*(1-np.mean(vals[8])/np.mean(vals[0])):.1f}")
        if 2 in ems:
            macro("abl" + key + "embtwo", f"{100*np.mean(ems[2]):.1f}")
        macro("abl" + key + "seeds", str(n))


_old_make_all6 = make_all
def make_all(macro):  # noqa: F811
    _old_make_all6(macro); table_ablations(macro)