"""Cross-model generalization summary (Phase 3).

Reads the per-model acceptance result JSONs (gpt-oss Phase-2 headline + any
Phase-3 models) and emits:
  * results/phase3_crossmodel.json  -- machine-readable comparison
  * paper/figures/crossmodel.{pdf,png} -- grouped bar of personal vs static MAT
  * a LaTeX-ready table printed to stdout

No model is re-run here; this only aggregates existing result files. Missing
model files are skipped (so it works whether or not Nemotron served).
"""
from __future__ import annotations

import json
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
RESULTS = ROOT / "results"
FIGS = ROOT / "figures"

# (display name, results json). row order = table order (gpt-oss first).
MODELS = [
    ("gpt-oss-120b", RESULTS / "phase2_accept_results.json"),
    ("gemma-4-31B-it", RESULTS / "phase3_gemma_accept_results.json"),
    ("Nemotron-3-Super-120B", RESULTS / "phase3_nemotron_accept_results.json"),
]


def _row(path: Path):
    d = json.loads(path.read_text())
    ow = d["overall_post_warmup"]
    per = d["summary"]  # per-session, for tail gap
    sessions = sorted(int(s) for s in per["personal_memory"].keys())
    last = str(sessions[-1])
    def mat(arm, block=ow, key=None):
        return block[arm]["MAT"] if key is None else block[arm][key]["MAT"]
    stat, pers, nomem = mat("static_global"), mat("personal_memory"), mat("no_memory")
    gap = 100.0 * (pers - stat) / stat
    tail_stat = per["static_global"][last]["MAT"]
    tail_pers = per["personal_memory"][last]["MAT"]
    tail_gap = 100.0 * (tail_pers - tail_stat) / tail_stat
    return {
        "no_memory": round(nomem, 2),
        "static_global": round(stat, 2),
        "personal_memory": round(pers, 2),
        "gap_pct": round(gap, 1),
        "tail_gap_pct": round(tail_gap, 1),
        "personal_seed_std": ow["personal_memory"].get("MAT_seed_std"),
        "n": ow["personal_memory"]["n"],
    }


def main():
    out = {}
    for name, path in MODELS:
        if path.exists():
            out[name] = _row(path)
            print(f"[ok] {name}: {out[name]}")
        else:
            print(f"[skip] {name}: {path.name} not found")
    (RESULTS / "phase3_crossmodel.json").write_text(json.dumps(out, indent=2))

    # LaTeX table body
    print("\n% --- LaTeX table rows (personal vs static vs none, +gap) ---")
    for name, r in out.items():
        print(f"{name} & {r['no_memory']:.2f} & {r['static_global']:.2f} & "
              f"{r['personal_memory']:.2f} & $+{r['gap_pct']:.0f}\\%$ & "
              f"$+{r['tail_gap_pct']:.0f}\\%$ \\\\")

    # grouped bar figure
    try:
        import matplotlib
        matplotlib.use("Agg")
        import matplotlib.pyplot as plt
        import numpy as np
        names = list(out.keys())
        x = np.arange(len(names))
        w = 0.26
        nomem = [out[n]["no_memory"] for n in names]
        stat = [out[n]["static_global"] for n in names]
        pers = [out[n]["personal_memory"] for n in names]
        fig, ax = plt.subplots(figsize=(7.2, 3.6))
        ax.bar(x - w, nomem, w, label="No memory", color="#9e9e9e")
        ax.bar(x, stat, w, label="Static datastore", color="#4C72B0")
        ax.bar(x + w, pers, w, label="Personal evicting (ours)", color="#C44E52")
        ymax = max(pers) * 1.30  # headroom for labels + legend
        ax.set_ylim(0, ymax)
        for xi, n in zip(x, names):
            ax.text(xi + w, out[n]["personal_memory"] + ymax * 0.015,
                    f"+{out[n]['gap_pct']:.0f}%", ha="center", fontsize=8,
                    color="#C44E52", fontweight="bold")
        ax.set_xticks(x)
        ax.set_xticklabels(names, fontsize=9)
        ax.set_ylabel("Mean accepted tokens (post-warmup)")
        ax.set_title("Personalized evicting memory generalizes across served models")
        ax.legend(fontsize=8, loc="upper center", ncol=3, frameon=False,
                  bbox_to_anchor=(0.5, 1.0))
        ax.grid(axis="y", alpha=0.3)
        fig.tight_layout()
        FIGS.mkdir(parents=True, exist_ok=True)
        fig.savefig(FIGS / "crossmodel.pdf")
        fig.savefig(FIGS / "crossmodel.png", dpi=150)
        print(f"\n[fig] wrote {FIGS/'crossmodel.pdf'}")
    except Exception as e:  # noqa: BLE001
        print(f"[fig] skipped: {e}")


if __name__ == "__main__":
    main()