#!/usr/bin/env python3
"""Generate LaTeX table fragments in paper/tables/ from registered run results (no computation).

Usage: python scripts/build_tables.py --graft <run_id> [--pilot <run_id>] [--tier1 <run_id>] [--confirmatory <run_id>] [--tier1-confirmatory <run_id>]
Each fragment starts with a comment naming its source run so that every number is traceable.
"""
from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

import numpy as np

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "src"))
OUT = ROOT / "paper" / "tables"

PRETTY = {
    "act_identity_all": "Activation graft, identity map", "act_procrustes_all": "Activation graft, Procrustes (task rows)",
    "act_unrelated_procrustes_all": "Activation graft, Procrustes (task-free rows)", "ctl_unrelated_data_map": "Activation graft, ridge (task-free rows)",
    "act_ridge_all": "Activation graft, ridge (task rows)", "act_cca_all": "Activation graft, CCA (task rows)",
    "ctl_sham_map": "Control: sham map (random orthogonal)", "ctl_norm_matched_random": "Control: norm-matched random rows",
    "ctl_shuffled_token": "Control: shuffled-item donor rows", "ctl_shuffled_layer": "Control: donor layer 3 rows (ridge)",
    "ctl_unrelated_donor": "Control: unrelated-skill donor (ridge)", "steer_a1.0": "Steering vector, $\\alpha=1$", "steer_a4.0": "Steering vector, $\\alpha=4$",
    "subspace_k8": "Subspace graft, $k=8$", "subspace_k32": "Subspace graft, $k=32$", "subspace_k128": "Subspace graft, $k=128$",
    "subspace_random_k32": "Control: random subspace, $k=32$", "adapter_graft": "Adapter graft (donor-free, ridge interface)",
    "baseline_supervised_adapter": "Baseline: supervised adapter (labels)", "baseline_distillation": "Baseline: distillation (donor outputs)",
    "weight_L0_mlp": "Weight graft: L0 MLP", "weight_L0_attn": "Weight graft: L0 attention", "weight_L0_block_ln": "Weight graft: L0 block",
    "weight_L012_blocks_embed": "Weight graft: L0--L2 blocks + $\\Omega$ embedding", "weight_embed_omega": "Weight graft: $\\Omega$ embedding only",
    "ctl_weight_shuffled_layer": "Control: donor L3 block into L0", "ctl_weight_unrelated_donor": "Control: unrelated-skill donor L0 block",
}


def fmt(x, nd=3):
    return "---" if x is None else f"{x:.{nd}f}"


def load(run):
    return json.loads((ROOT / "experiments" / "summaries" / run / "results.json").read_text())


def pairs_of(r, same=None):
    pm = r["pair_meta"]
    return [p for p, m in pm.items() if same is None or m["same_init_seed"] == same]


def mean_sd(r, cond, pairs, split="target_iid", key="accuracy"):
    v = [r["conditions"][cond][p][split][key] for p in pairs if p in r["conditions"].get(cond, {}) and split in r["conditions"][cond][p]]
    if not v:
        return None, None, 0
    return float(np.mean(v)), float(np.std(v)), len(v)


def table_main(r, run, name="tab_main_effects.tex"):
    same, diff = pairs_of(r, True), pairs_of(r, False)
    rows = []
    for cond in PRETTY:
        if cond not in r["conditions"]:
            continue
        ts, tsd, n1 = mean_sd(r, cond, same)
        cs, _, _ = mean_sd(r, cond, same, "control_iid")
        td, tdd, n2 = mean_sd(r, cond, diff)
        cd, _, _ = mean_sd(r, cond, diff, "control_iid")
        # control-adjusted effect vs sham on the matched pairs (mean of per-pair estimates and CI bounds)
        adj = [r["contrasts"][p][cond]["control_adjusted"] for p in same + diff if cond in r["contrasts"].get(p, {}) and "control_adjusted" in r["contrasts"][p][cond]]
        adj_s = f"{np.mean([a['estimate'] for a in adj]):+.3f} [{np.mean([a['low'] for a in adj]):+.3f}, {np.mean([a['high'] for a in adj]):+.3f}]" if adj else "---"
        rows.append(f"{PRETTY[cond]} & {fmt(ts)} ({fmt(tsd,2)}) & {fmt(cs)} & {fmt(td) if n2 else '---'} {('('+fmt(tdd,2)+')') if n2 else ''} & {fmt(cd) if n2 else '---'} & {adj_s} \\\\")
    base_t = [r["conditions"]["base"][k]["target_iid"]["accuracy"] for k in r["conditions"]["base"] if k.startswith("recipient_s") and "_s2_" not in k and "w96" not in k and "L6" not in k]
    hdr = f"""% source: experiments/summaries/{run}/results.json (generated by scripts/build_tables.py)
\\begin{{tabular}}{{lccccc}}
\\toprule
 & \\multicolumn{{2}}{{c}}{{same initialization ($n={len(same)}$ pairs)}} & \\multicolumn{{2}}{{c}}{{different initialization ($n={len(diff)}$)}} & control-adjusted \\\\
\\cmidrule(lr){{2-3}} \\cmidrule(lr){{4-5}}
Condition & target (sd) & control & target (sd) & control & $\\Delta_T^{{\\mathrm{{adj}}}}$ [95\\% CI] \\\\
\\midrule
Recipient, unmodified & {fmt(np.mean(base_t))} ({fmt(np.std(base_t),2)}) & 1.000 & & & \\\\
"""
    (OUT / name).write_text(hdr + "\n".join(rows) + "\n\\bottomrule\n\\end{tabular}\n")


def table_maps(r, run):
    same, diff = pairs_of(r, True), pairs_of(r, False)
    rows = []
    for fam, label in (("identity", "identity"), ("procrustes", "Procrustes (task rows)"), ("unrelated_procrustes", "Procrustes (task-free rows)"), ("ridge", "ridge (task rows)"), ("unrelated_data", "ridge (task-free rows)"), ("cca", "CCA (task rows)"), ("random_orth", "random orthogonal (sham)")):
        def get(p):
            m = r["maps"][f"{p}|{fam}"]
            return m.get("r2_fit", m.get("r2_on_target_rows"))
        a = [get(p) for p in same]
        b = [get(p) for p in diff]
        npar = r["maps"][f"{same[0]}|{fam}"]["n_params"]
        rows.append(f"{label} & {npar} & {fmt(np.mean(a))} & {fmt(np.mean(b))} \\\\")
    txt = f"""% source: experiments/summaries/{run}/results.json
\\begin{{tabular}}{{lrcc}}
\\toprule
Map & parameters & $R^2$ on $\\Omega$ rows, same init & $R^2$, different init \\\\
\\midrule
""" + "\n".join(rows) + "\n\\bottomrule\n\\end{tabular}\n"
    (OUT / "tab_maps.tex").write_text(txt)


def table_mismatch(r, run):
    rows = []
    for cond, d in r["conditions"].items():
        if cond.startswith("mismatch_") or cond.startswith("stride2_"):
            for p, v in d.items():
                rows.append(f"{p.replace('_', '\\_')} & {cond.replace('_', '\\_')} & {fmt(v['target_iid']['accuracy'])} & {fmt(v['target_iid'].get('omega_step_acc'))} & {fmt(v['control_iid']['accuracy'])} & {fmt(v.get('r2_fit'))} \\\\")
    for key in ("recipient_w96_s0", "recipient_L6_s0", "recipient_s2_s0", "donor_s2_s0"):
        if key in r["conditions"]["base"]:
            v = r["conditions"]["base"][key]
            rows.insert(0, f"{key.replace('_', '\\_')} & base & {fmt(v['target_iid']['accuracy'])} & {fmt(v['target_iid'].get('omega_step_acc'))} & {fmt(v['control_iid']['accuracy'])} & --- \\\\")
    txt = f"""% source: experiments/summaries/{run}/results.json
\\begin{{tabular}}{{llcccc}}
\\toprule
Pair & condition & target & $\\Omega$-step & control & map $R^2$ \\\\
\\midrule
""" + "\n".join(rows) + "\n\\bottomrule\n\\end{tabular}\n"
    (OUT / "tab_mismatch.tex").write_text(txt)


def table_mechanism(r, run):
    rows = []
    for cond in ("mech_donor_ablated_attn0", "mech_rescue_from_intact_donor", "mech_rescue_sham", "mech_reverse_recipient_to_donor"):
        for p, v in r["conditions"].get(cond, {}).items():
            rows.append(f"{cond.replace('mech_', '').replace('_', ' ')} & {p.replace('_', '\\_')} & {fmt(v['target_iid']['accuracy'])} & {fmt(v['target_iid'].get('omega_step_acc'))} & {fmt(v['control_iid']['accuracy'])} \\\\")
    h5 = r["rule_vs_cache"]
    for k, v in h5.items():
        if isinstance(v, dict) and "seen" in v:
            rows.append(f"rule vs cache: {k.replace('h5_', '').replace('_', ' ')} & seen / unseen $\\Omega$ inputs & {fmt(v['seen']['estimate'])} [{fmt(v['seen']['low'])}, {fmt(v['seen']['high'])}] & {fmt(v['unseen']['estimate'])} [{fmt(v['unseen']['low'])}, {fmt(v['unseen']['high'])}] & --- \\\\")
    txt = f"""% source: experiments/summaries/{run}/results.json
\\begin{{tabular}}{{llccc}}
\\toprule
Test & pair / stratum & target & $\\Omega$-step & control \\\\
\\midrule
""" + "\n".join(rows) + "\n\\bottomrule\n\\end{tabular}\n"
    (OUT / "tab_mechanism.tex").write_text(txt)


def table_tier1(r, run, name="tab_tier1.tex"):
    rows = []
    for p, pr in r["pairs"].items():
        for tn, tr in pr["tasks"].items():
            conds = tr["conditions"]
            best_name, best = max(((k, v) for k, v in conds.items() if not k.startswith("ctl") and "random_orth" not in k), key=lambda kv: kv[1]["accuracy"])
            ctl_name, ctl = max(((k, v) for k, v in conds.items() if k.startswith("ctl") or "random_orth" in k), key=lambda kv: kv[1]["accuracy"])
            g = best["gain_vs_zero"]
            rows.append(f"{p.replace('_', '\\_')} & {tn} & {fmt(tr['recipient_zero'],2)} & {fmt(tr['recipient_few'],2)} & {fmt(tr['recipient_own_vector'],2)} & {fmt(best['accuracy'],2)} & {best_name.replace('_', '\\_').replace('|', ' ')} & {fmt(ctl['accuracy'],2)} & {g['estimate']:+.2f} [{g['low']:+.2f}, {g['high']:+.2f}] \\\\")
    txt = f"""% source: experiments/summaries/{run}/results.json
\\begin{{tabular}}{{llccccccc}}
\\toprule
Pair & task & zero-shot & few-shot & own vector & best transfer & condition & best control & gain vs zero [95\\% CI] \\\\
\\midrule
""" + "\n".join(rows) + "\n\\bottomrule\n\\end{tabular}\n"
    (OUT / name).write_text(txt)


def table_tier1_localization(r, run, name="tab_tier1_localization.tex"):
    rows = []
    for mk, m in r["models"].items():
        loc = m.get("localization")
        if not loc:
            continue
        for tn, t in loc["tasks"].items():
            rows.append(f"{mk} ({m['n_layers']} layers, $d={m['hidden']}$) & {tn} & {fmt(t['zero_shot'],2)} & {fmt(t['few_shot'],2)} & L{t['best_layer']} & {fmt(t['sufficiency_by_layer'][str(t['best_layer'])],2)} & {fmt(t['random_vector_same_norm_at_best'],2)} & {fmt(t['necessity_fewshot_minus_direction'],2)} \\\\")
    txt = f"""% source: experiments/summaries/{run}/results.json
\\begin{{tabular}}{{llcccccc}}
\\toprule
Model & task & zero-shot & few-shot & best layer & vector added (zero-shot) & random vector & few-shot minus direction \\\\
\\midrule
""" + "\n".join(rows) + "\n\\bottomrule\n\\end{tabular}\n"
    (OUT / name).write_text(txt)


def table_tier1_confirmatory(run):
    c = json.loads((ROOT / "experiments" / "summaries" / run / "tier1_confirmatory.json").read_text())
    rows = []
    for p, pr in c["pairs"].items():
        for tn, t in pr["tasks"].items():
            if "error" in t:
                continue
            el = "yes" if t["eligible"] else "no"
            rows.append(f"{p.replace('->', ' $\\to$ ')} & {tn} & {el} & {t['selected'].replace('_', '\\_').replace('|', ' ')} & {fmt(t['conf_zero'],2)} & {fmt(t['conf_few'],2)} & {fmt(t['conf_own'],2)} & {fmt(t['conf_selected'],2)} & {fmt(t['best_control_acc'],2)} & {t['diff']:+.2f} [{t['ci'][0]:+.2f}, {t['ci'][1]:+.2f}]{' $\\checkmark$' if t['supported'] else ''} & {fmt(max(t['shuffled_query_diagnostic'].values()) if t['shuffled_query_diagnostic'] else None,2)} \\\\")
        rows.append(f"\\multicolumn{{11}}{{l}}{{\\emph{{{p.replace('->', ' $\\to$ ')}: H6 {pr['H6']} ({pr['n_supported_tasks']} task(s))}}}} \\\\ \\addlinespace")
    txt = f"""% source: experiments/summaries/{run}/tier1_confirmatory.json (scripts/tier1_analysis.py)
\\begin{{tabular}}{{llllcccccll}}
\\toprule
Pair & task & eligible & selected condition (dev) & zero & few & own & selected & best null control & selected $-$ control [95\\% CI] & shuffled-query diag. \\\\
\\midrule
""" + "\n".join(rows) + "\n\\bottomrule\n\\end{tabular}\n"
    (OUT / f"tab_tier1_confirmatory_{run.split('_')[-1][:6]}.tex").write_text(txt)


def table_pilot(r, run):
    rows = []
    for en in ("target_iid", "control_iid", "comp_heldout", "len_ood"):
        d, rc = r["across_seeds"][en]["donor"], r["across_seeds"][en]["recipient"]
        g = r["seeds"]["0"]["gap"][en]
        rows.append(f"{en.replace('_', '\\_')} & {fmt(r['chance_level'][en])} & {fmt(d['mean'])} ({fmt(d['sd'] or 0,3)}) & {fmt(rc['mean'])} ({fmt(rc['sd'] or 0,3)}) & {g['estimate']:+.3f} [{g['low']:+.3f}, {g['high']:+.3f}] \\\\")
    txt = f"""% source: experiments/summaries/{run}/results.json
\\begin{{tabular}}{{lcccc}}
\\toprule
Stratum & chance & donor (sd) & recipient (sd) & paired gap, seed 0 [95\\% CI] \\\\
\\midrule
""" + "\n".join(rows) + "\n\\bottomrule\n\\end{tabular}\n"
    (OUT / "tab_pilot_gap.tex").write_text(txt)


def table_models(index_path: Path):
    idx = json.loads(index_path.read_text())
    rows = []
    for k, v in idx.items():
        m = v["model"]
        t = v["task"]
        rows.append(f"{k.replace('_', '\\_')} & {m['n_layers']} & {m['d_model']} & {m['n_heads']} & {', '.join(t['excluded_operators']) or '---'} & {'vex' if 'vex' in t['operators'] else '---'} & {t.get('trace_stride', 1)} & {fmt(v['final_dev_acc'])} & {v['train_hash']} \\\\")
    txt = f"""% source: {index_path.relative_to(ROOT)}
\\begin{{tabular}}{{lcccccccl}}
\\toprule
Model & layers & $d$ & heads & excluded & extra op. & stride & dev acc. & train split hash \\\\
\\midrule
""" + "\n".join(rows) + "\n\\bottomrule\n\\end{tabular}\n"
    (OUT / "tab_models.tex").write_text(txt)


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--graft")
    ap.add_argument("--pilot")
    ap.add_argument("--tier1", nargs="*", default=[])
    ap.add_argument("--confirmatory")
    ap.add_argument("--tier1-confirmatory", nargs="*", default=[])
    args = ap.parse_args()
    OUT.mkdir(parents=True, exist_ok=True)
    if args.graft:
        r = load(args.graft)
        table_main(r, args.graft)
        table_maps(r, args.graft)
        table_mismatch(r, args.graft)
        table_mechanism(r, args.graft)
    if args.confirmatory:
        table_main(load(args.confirmatory), args.confirmatory, "tab_confirmatory.tex")
    if args.pilot:
        table_pilot(load(args.pilot), args.pilot)
    for run in args.tier1:
        r = load(run)
        suf = run.split("_")[-1][:6]
        table_tier1(r, run, f"tab_tier1_{suf}.tex")
        table_tier1_localization(r, run, f"tab_tier1_localization_{suf}.tex")
    for run in args.tier1_confirmatory:
        table_tier1_confirmatory(run)
    idx = ROOT / "artifacts" / "models" / "matrix_v1" / "index.json"
    if idx.exists():
        table_models(idx)
    print("tables written:", sorted(p.name for p in OUT.glob("*.tex")))


if __name__ == "__main__":
    main()