#!/usr/bin/env python
"""Generate LaTeX table fragments and the numbers macro file from structured results (never typed by hand).
Usage: python scripts/make_tables.py  -> paper/tables/*.tex, paper/tables/numbers.tex
"""
from __future__ import annotations

import json, re, sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
OUT = Path("paper/tables"); OUT.mkdir(parents=True, exist_ok=True)
NUM: dict[str, str] = {}


def macro(name: str, value) -> None:
    name = re.sub(r"[^A-Za-z]", "", name)
    NUM[name] = value if isinstance(value, str) else f"{value}"


def fmt_int(n) -> str:
    return f"{int(n):,}"


def table_configs():
    pc = json.loads(Path("results/tables/param_counts.json").read_text())
    rows = []
    order = ["prisme_103m", "prisme_459m", "prisme_1180m"]
    labels = {"prisme_103m": "PRISME-103M", "prisme_459m": "PRISME-459M", "prisme_1180m": "PRISME-1.18B"}
    head = r"""\begin{table}[t]\centering\small
\caption{Dimensioned configurations and exact unique parameter counts (meta-device instantiation,
\texttt{scripts/count\_params.py}). Backbone counts equal the reference design exactly; workspace+controller counts
differ from the proposal's figures (reconciled in Appendix~\ref{app:params}).}
\label{tab:configs}
\begin{tabular}{lrrr}\toprule
Property & PRISME-103M & PRISME-459M & PRISME-1.18B\\\midrule
"""
    cfgs = {}
    for k in order:
        c = json.loads(Path(f"configs/{k}.yaml").read_text().split("model:")[1].replace("\n", ",").replace("  ", "").replace(":", '":').replace(",", ',"').replace('{"', "{").replace(",}", "}") if False else "{}")
    import yaml
    cfgs = {k: yaml.safe_load(open(f"configs/{k}.yaml"))["model"] for k in order}
    def row(name, f):
        return name + " & " + " & ".join(f(k) for k in order) + r"\\" + "\n"
    body = ""
    body += row("Backbone layers $L$", lambda k: str(cfgs[k]["n_layers"]))
    body += row("Hidden $d$ / FFN", lambda k: f"{cfgs[k]['d_model']} / {cfgs[k]['ffn_width']}")
    body += row("Query / KV heads", lambda k: f"{cfgs[k]['n_heads']} / {cfgs[k]['n_kv_heads']}")
    body += row("Workspace $w$ / FFN$_w$", lambda k: f"{cfgs[k]['w']} / {cfgs[k]['ffn_w']}")
    body += row("Slots $M$, blocks $L_w$, cap $R_{\\max}$", lambda k: f"{cfgs[k]['n_slots']}, {cfgs[k]['n_editor_blocks']}, {cfgs[k]['r_max']}")
    body += r"\midrule" + "\n"
    body += row("Backbone parameters", lambda k: fmt_int(pc[k]["backbone"]))
    body += row("Workspace + controller", lambda k: fmt_int(pc[k]["workspace_controller"]))
    body += row("\\textbf{Total}", lambda k: "\\textbf{" + fmt_int(pc[k]["total"]) + "}")
    body += row("Proposal's workspace figure", lambda k: fmt_int(pc[k]["reference_brief"]["workspace"]))
    tail = r"\bottomrule\end{tabular}\end{table}" + "\n"
    (OUT / "configs.tex").write_text(head + body + tail)
    for k in order:
        macro("params" + {"prisme_103m": "Small", "prisme_459m": "Medium", "prisme_1180m": "Large"}[k] + "Total", fmt_int(pc[k]["total"]))
        macro("params" + {"prisme_103m": "Small", "prisme_459m": "Medium", "prisme_1180m": "Large"}[k] + "Ws", fmt_int(pc[k]["workspace_controller"]))


def table_ops():
    o = json.loads(Path("results/tables/ops_459m_T4096.json").read_text())
    names = {"backbone_only": "Backbone only", "4_actions_k1": "4 actions, 1 slot each", "4_actions_k2": "4 actions, 2 slots each",
             "8_actions_k2": "8 actions, 2 slots each", "8_actions_dense16": "8 actions, dense 16-slot writes"}
    head = r"""\begin{table}[t]\centering\small
\caption{Analytic operations per generated token, PRISME-459M, context $T=4096$ (\texttt{scripts/estimate\_flops.py}).
``Proposal'' is the reference table of the original brief; ``implemented'' is the accounting of
Section~\ref{sec:compute} with the controller evaluated once per executed action plus once to stop. These are
arithmetic counts, not measured timings.}
\label{tab:ops}
\begin{tabular}{lrrrrr}\toprule
Policy & Proposal ($10^9$) & Implemented ($10^9$) & editor & controller & readout\\\midrule
"""
    body = ""
    for k, lab in names.items():
        r = o["rows"][k]
        ed = (r["editor_source_attn"] + r["editor_slot_attn"] + r["editor_ffn"] + r["editor_gate"]) / 1e9
        body += f"{lab} & {o['brief_reference_459m_T4096'][k]/1e9:.3f} & {r['total']/1e9:.3f} & {ed:.3f} & {r['controller']/1e9:.3f} & {r['readout']/1e9:.3f}\\\\\n"
    (OUT / "ops.tex").write_text(head + body + r"\bottomrule\end{tabular}\end{table}" + "\n")
    o103 = json.loads(Path("results/tables/ops_103m_T512.json").read_text())
    macro("costsingle", f"{o103['action_cost_proxy'][0]:.3f}"); macro("costpair", f"{o103['action_cost_proxy'][-1]:.3f}")


def table_families():
    man = json.loads(Path("data/splits/synthetic_v1/manifest.json").read_text())
    head = r"""\begin{table}[t]\centering\small
\caption{Synthetic suite v1: splits fixed before sampling. $n$ = items after cross-split exact-duplicate removal;
``dropped'' = duplicates removed. Test splits are locked until pre-registration is frozen.}
\label{tab:families}
\begin{tabular}{llllrr}\toprule
Split & Families & Templates & Depths & $n$ & dropped\\\midrule
"""
    body = ""
    for name, m in man.items():
        s = m["spec"]
        fam = ", ".join(s["families"]) if len(s["families"]) > 1 else s["families"][0]
        if len(s["families"]) == 5:
            fam = "arith, bind, fsm, excl, sched"
        body += f"{name.replace('_', '\\_')} & {fam} & {','.join(map(str, s['templates']))} & {','.join(map(str, s['depths']))} & {m['n']:,} & {m['duplicates_dropped']}\\\\\n"
    (OUT / "families.tex").write_text(head + body + r"\bottomrule\end{tabular}\end{table}" + "\n")
    macro("ntrainsynth", f"{man['train']['n']:,}")


def write_numbers():
    lines = ["% AUTO-GENERATED by scripts/make_tables.py — do not edit"]
    for k, v in sorted(NUM.items()):
        lines.append(f"\\newcommand{{\\{k}}}{{{v}}}")
    (OUT / "numbers.tex").write_text("\n".join(lines) + "\n")


if __name__ == "__main__":
    table_configs(); table_ops(); table_families()
    try:
        from scripts.make_results_tables import make_all  # results-dependent tables (added later)
        make_all(macro)
    except Exception as e:  # noqa: BLE001
        print("results tables skipped:", e)
    write_numbers()
    print("tables written:", sorted(p.name for p in OUT.glob("*.tex")))