#!/usr/bin/env python
"""Regenerate every figure of the manuscript from structured results. Figures whose inputs are absent are skipped
and reported. Usage: python scripts/reproduce_figures.py --results results/raw/ [--runs experiments/runs]
"""
from __future__ import annotations

import argparse, glob, json, sys
from collections import defaultdict
from pathlib import Path

import numpy as np

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from scripts.plotstyle import CAT, SEQ_BLUE, setup, save  # noqa: E402

import matplotlib.pyplot as plt  # noqa: E402

MADE: list[str] = []
SKIPPED: list[str] = []


def load_metrics(run_dir):
    p = Path(run_dir) / "metrics.jsonl"
    return [json.loads(l) for l in open(p)] if p.exists() else []


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


# ---------------------------------------------------------------- 1. learnability: val loss by budget vs step
def fig_learnability(runs_root, prefixes, name, title):
    panels = [(p, runs_matching(runs_root, p)) for p in prefixes]
    panels = [(p, r) for p, r in panels if r]
    if not panels:
        SKIPPED.append(name); return
    fig, axes = plt.subplots(1, len(panels), figsize=(2.6 * len(panels), 2.4), sharey=True, squeeze=False)
    for ax, (p, runs) in zip(axes[0], panels):
        for run in runs[:1]:
            rows = [r for r in load_metrics(run) if any(k.startswith("val_loss_b") for k in r)]
            budgets = sorted({int(k[10:]) for r in rows for k in r if k.startswith("val_loss_b")})
            for i, b in enumerate(budgets):
                xs = [r["step"] + 1 for r in rows if f"val_loss_b{b}" in r]; ys = [r[f"val_loss_b{b}"] for r in rows if f"val_loss_b{b}" in r]
                ax.plot(xs, ys, marker="o", color=SEQ_BLUE[min(i + 1, 6)] if b else "#52514e", label=f"budget {b}")
        ax.set_title(p.replace("_", " "), loc="left"); ax.set_xlabel("training step")
    axes[0][0].set_ylabel("held-out answer-span loss (nats)"); axes[0][-1].legend(loc="upper right")
    fig.suptitle(title, x=0.01, ha="left", fontsize=8.5)
    MADE.append(save(fig, name))


# ---------------------------------------------------------------- 2. labels: gain distribution + variance decomposition
def fig_labels(analyzed_glob):
    files = sorted(glob.glob(analyzed_glob))
    if not files:
        SKIPPED.append("labels"); return
    d = json.loads(Path(files[-1]).read_text())
    part = d.get("calib") or d.get("train")
    aud = part.get("audit", {})
    fig, axes = plt.subplots(1, 3, figsize=(9.0, 2.5), gridspec_kw={"wspace": 0.45})
    ax = axes[0]
    depths = sorted(int(k) for k in part["by_prefix_depth"])
    ax.bar([str(k) for k in depths], [part["by_prefix_depth"][str(k)]["frac_pos"] for k in depths], color=CAT[0], width=0.6)
    ax.axhline(0.5, color="#52514e", lw=0.8, ls="--"); ax.set_ylim(0, 0.7)
    ax.set_xlabel("prefix depth r (edits already applied)"); ax.set_ylabel("fraction of gains > 0"); ax.set_title("Sign of logged gains", loc="left")
    ax = axes[1]
    labels = ["between states", "within state (action)"]
    vals = [aud.get("var_between_states", 0), aud.get("var_within_state_actions", 0)]
    tot = sum(vals) or 1
    ax.bar(labels, [v / tot for v in vals], color=[CAT[0], CAT[1]], width=0.6)
    for i, v in enumerate(vals):
        ax.text(i, v / tot + 0.02, f"{100*v/tot:.0f}%", ha="center", fontsize=7)
    ax.set_ylim(0, 1.05); ax.set_ylabel("share of gain variance"); ax.set_title("Variance of audited gains", loc="left")
    ax = axes[2]
    names = ["random\naction", "perfect stop,\nrandom action", "one-step\noracle"]
    vals = [aud.get("random_action_mean", 0), aud.get("oracle_state_level_mean", 0), aud.get("oracle_one_step_mean", 0)]
    ax.bar(names, vals, color=[CAT[2], CAT[3], CAT[0]], width=0.6)
    for i, v in enumerate(vals):
        ax.text(i, v + 0.0005, f"{v:.4f}", ha="center", fontsize=7)
    ax.set_ylabel("mean gain per state (nats)"); ax.set_title("Oracle decomposition", loc="left")
    MADE.append(save(fig, "labels_analysis"))


# ---------------------------------------------------------------- 3. routing Pareto: loss vs extra ops
def fig_routing(results_dir):
    files = sorted(glob.glob(f"{results_dir}/routing_*.json"))
    if not files:
        SKIPPED.append("routing_pareto"); return
    for f in files:
        d = json.loads(Path(f).read_text()); recs = d["records"]
        fig, ax = plt.subplots(figsize=(4.2, 2.8))
        groups = defaultdict(list)
        for r in recs:
            if r.get("policy") in ("oracle_two_step",):
                continue
            groups[r["policy"]].append(r)
        order = ["budget0", "random_forced", "uniform", "fixed_pair", "entropy", "norm", "controller", "controller_locked", "oracle1_greedy"]
        for i, pol in enumerate([p for p in order if p in groups]):
            rs = sorted(groups[pol], key=lambda r: r["extra_ops_ratio"])
            xs = [r["extra_ops_ratio"] for r in rs]; ys = [r["loss"] for r in rs]
            style = dict(marker="o", color=CAT[i % 8], ls="-" if len(rs) > 1 else "none")
            if pol == "oracle1_greedy":
                style.update(ls="--", color="#52514e")
            if pol == "controller_locked":
                style.update(marker="*", markersize=9, ls="none")
            ax.plot(xs, ys, label=pol.replace("_", " ") + (" (diagnostic)" if "oracle" in pol else ""), **style)
        ax.set_xlabel("extra operations per token (× backbone)"); ax.set_ylabel("held-out loss (nats)")
        ax.set_title(f"Quality vs analytic cost — {d['split']}", loc="left"); ax.legend(ncol=2)
        MADE.append(save(fig, "routing_" + Path(f).stem.replace("routing_", "")))


# ---------------------------------------------------------------- 4. latency
def fig_latency(results_dir):
    files = sorted(glob.glob(f"{results_dir}/latency_*.json"))
    if not files:
        SKIPPED.append("latency"); return
    for f in files:
        d = json.loads(Path(f).read_text()); recs = [r for r in d["records"] if r["batch"] == 1]
        fig, ax = plt.subplots(figsize=(4.2, 2.6))
        groups = defaultdict(list)
        for r in recs:
            groups[r["policy"]].append(r)
        for i, (pol, rs) in enumerate(groups.items()):
            rs = sorted(rs, key=lambda r: r["budget"])
            ax.errorbar([r["budget"] for r in rs], [r["median_ms"] for r in rs], yerr=[[0] * len(rs), [r["p95_ms"] - r["median_ms"] for r in rs]],
                        marker="o", color=CAT[i % 8], label=pol.replace("_", " "), capsize=2, lw=1.2)
        ax.set_xlabel("executed edits per token (budget)"); ax.set_ylabel("ms per token (median, +p95)")
        ax.set_title(f"Batch-1 decode latency — {d['device']}", loc="left"); ax.legend()
        MADE.append(save(fig, Path(f).stem))


# ---------------------------------------------------------------- 5. exact match by regime/policy
def fig_exact_match(results_dir):
    files = sorted(glob.glob(f"{results_dir}/eval_*_dev_*.json")) + sorted(glob.glob(f"{results_dir}/eval_*_test_*.json"))
    files = [f for f in files if "exact_match" in Path(f).read_text()[:400] or True]
    data = defaultdict(dict)
    for f in files:
        d = json.loads(Path(f).read_text())
        if "exact_match" not in d:
            continue
        for pol, em in d["exact_match"].items():
            data[(d["system"], pol)][d["split"]] = em["accuracy"]
    if not data:
        SKIPPED.append("exact_match"); return
    splits = sorted({s for v in data.values() for s in v})
    fig, ax = plt.subplots(figsize=(6.4, 2.6))
    keys = sorted(data)
    width = 0.8 / max(len(keys), 1)
    for i, k in enumerate(keys):
        xs = np.arange(len(splits)) + i * width
        ax.bar(xs, [data[k].get(s, np.nan) for s in splits], width=width * 0.95, color=CAT[i % 8], label=f"{k[0]} / {k[1]}")
    ax.set_xticks(np.arange(len(splits)) + 0.4 - width / 2); ax.set_xticklabels([s.replace("_", "\n") for s in splits])
    ax.set_ylabel("exact match"); ax.set_title("Free-generation exact match by regime", loc="left"); ax.legend(ncol=2, fontsize=6)
    MADE.append(save(fig, "exact_match_regimes"))


# ---------------------------------------------------------------- 6. pretraining curves
def fig_pretrain(runs_root):
    runs = runs_matching(runs_root, "pretrain103m") + runs_matching(runs_root, "pretrain_enlarged102m")
    if not runs:
        SKIPPED.append("pretrain"); return
    fig, ax = plt.subplots(figsize=(4.2, 2.6))
    for i, run in enumerate(runs):
        rows = load_metrics(run)
        tr = [(r["tokens"], r["loss"]) for r in rows if "loss" in r]; va = [(r["tokens"], r["val_loss"]) for r in rows if "val_loss" in r]
        name = Path(run).name.split("_seed")[0]
        if tr:
            ax.plot([t / 1e6 for t, _ in tr], [l for _, l in tr], color=CAT[i], alpha=0.35, lw=0.8)
        if va:
            ax.plot([t / 1e6 for t, _ in va], [l for _, l in va], color=CAT[i], marker="o", label=f"{name} (val)")
    ax.set_xlabel("training tokens (M)"); ax.set_ylabel("cross-entropy (nats)"); ax.set_title("Stage A pretraining (FineWeb-Edu)", loc="left"); ax.legend()
    MADE.append(save(fig, "pretrain_curves"))


def main(argv=None):
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--results", default="results/raw/"); ap.add_argument("--runs", default="experiments/runs"); ap.add_argument("--analyzed", default="results/analyzed")
    args = ap.parse_args(argv); setup()
    fig_learnability(args.runs, ["p4_editor_v1_scaled", "p4_editor_v1_Ustd02", "p4_editor_v1_Ustd05"], "learnability_bootstrap", "Stage B bootstrap on the frozen tiny backbone (dev-iid, random forced actions)")
    fig_learnability(args.runs, ["p4v2_prisme", "p4v2_gr", "p4v2_dense"], "learnability_v2", "Stage B on the v2 tiny backbone: sparse editor vs global recurrence vs dense writes")
    fig_labels(f"{args.analyzed}/labels_*.json")
    fig_routing(args.results); fig_latency(args.results); fig_exact_match(args.results); fig_pretrain(args.runs)
    print("made:", MADE); print("skipped (inputs absent):", SKIPPED)


if __name__ == "__main__":
    main()