"""Final paper figures (CPU only), regenerated from consolidated results.

  fig_placement.png : band-entrance cliff + structural nulls
  fig_ladder.png    : MBPP-E2B attribution ladder, hard bucket, Wilson CIs
  fig_scale.png     : cross-scale / cross-task attribution grid
  fig_transfer.png  : substrate transfer panel (HumanEval / Rust / BW)
"""

import json
import math
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

ROOT = Path(__file__).resolve().parent.parent
OUT = ROOT / "results-loop"

BLUE, GREEN, ORANGE, GRAY, RED = ("#2b6cb0", "#2f855a", "#dd6b20",
                                  "#8a8f98", "#c53030")


def wilson(c, n, z=1.96):
    p = c / n
    d = 1 + z * z / n
    ctr = (p + z * z / (2 * n)) / d
    hw = z * math.sqrt(p * (1 - p) / n + z * z / (4 * n * n)) / d
    return ctr - hw, ctr + hw


def ev(path, k):
    d = json.load(open(path))
    v = d["ks"][str(k)]
    return v["acc"], v["by_label"].get("hard", float("nan"))


def best_hard(path, exclude0=True):
    d = json.load(open(path))
    items = [(int(k), v) for k, v in d["ks"].items()
             if not (exclude0 and k == "0")]
    k, v = max(items, key=lambda kv: kv[1]["by_label"].get("hard", 0))
    return k, v["acc"], v["by_label"].get("hard", 0)


def style(ax):
    ax.grid(True, color="#e8e8e8", lw=0.7)
    ax.set_axisbelow(True)
    for s in ("top", "right"):
        ax.spines[s].set_visible(False)


# ---------------- fig 1: placement ----------------
def fig_placement():
    entr = []  # (entrance_layer, overall@bestk, hard@bestk)
    for lo in (9, 11, 12, 13, 14):
        f = ROOT / f"results-band-{lo}_30" / f"eval_code_band_{lo}_30.json"
        if lo == 14:
            f = OUT / "eval_code_code_s0_full.json"
        k, a, h = best_hard(f)
        entr.append((lo, a, h))
    base_a, base_h = ev(OUT / "eval_code_code_s0_full.json", 0)

    fig, ax = plt.subplots(figsize=(6.4, 4.2))
    style(ax)
    xs = [e[0] for e in entr]
    ax.plot(xs, [e[2] for e in entr], "-s", color=BLUE, lw=2, ms=6,
            label="hard bucket (best k)")
    ax.plot(xs, [e[1] for e in entr], "-o", color=GRAY, lw=2, ms=5,
            label="overall (same k)")
    ax.axhline(base_h, color=BLUE, lw=1, ls=":", alpha=0.6)
    ax.axhline(base_a, color=GRAY, lw=1, ls=":", alpha=0.6)
    ax.annotate("base hard", (9.1, base_h + 0.012), fontsize=8, color=BLUE)
    ax.annotate("base overall", (9.1, base_a + 0.012), fontsize=8, color=GRAY)
    ax.axvspan(13.5, 14.5, color="#ebf4ff", zorder=0)
    ax.annotate("lens-identified\nworkspace entrance", (13.55, 0.60),
                fontsize=8, color=BLUE)
    ax.set_xticks(xs)
    ax.set_xlabel("loop entrance layer (exit fixed at L30)")
    ax.set_ylabel("MBPP pass@1")
    ax.set_title("Placement cliff: the retrofit works only at the "
                 "lens boundary (L14)", fontsize=11)
    ax.legend(fontsize=8, frameon=False, loc="center left")
    fig.text(0.13, 0.005,
             "Entrances 17/24 (not shown): structurally null — KV sharing "
             "makes k>0 bit-identical to k=0.", fontsize=7.5, color="#666")
    fig.tight_layout(rect=(0, 0.03, 1, 1))
    fig.savefig(OUT / "fig_placement.png", dpi=140, facecolor="white",
                bbox_inches="tight")
    print("wrote fig_placement.png")


# ---------------- fig 2: attribution ladder ----------------
def fig_ladder():
    NH = 55
    rows = []  # (label, hard_acc, n, color, note)

    def add(label, h, n, color, note=""):
        rows.append((label, h, n, color, note))

    _, h0 = ev(OUT / "eval_code_code_s0_full.json", 0)
    add("base (k=0, exact)", h0, NH, GRAY)
    _, hu = ev(OUT / "eval_code_untrained_full.json", 4)
    add("untrained loop (perturbation floor)", hu, NH, GRAY)
    _, hf = ev(OUT / "eval_code_code_ff_full.json", 1)
    add("trained FF (no recurrence)", hf, NH, ORANGE)
    _, hp = ev(OUT / "eval_code_pause16.json", 1)
    add("pause-16 registers (width)", hp, NH, ORANGE)
    _, hl = ev(OUT / "eval_code_code_s0_full.json", 4)
    add("loop k=4 (depth)", hl, NH, BLUE, "seed mean 0.375 ± 0.055")
    d = json.load(open(OUT / "eval_rung2_s1.json"))
    add("rung-2: + band LoRA (k=4)", d["4"]["by_label"]["hard"], 28, BLUE)
    _, hd = ev(OUT / "eval_code_distill_s1.json", 1)
    add("plan-distilled FF", hd, NH, GREEN, "8-run mean 0.457 ± 0.046")
    d = json.load(open(OUT / "eval_budgetcot.json"))
    add("budget-CoT (50 visible tok)", d["by_label"]["hard"], NH, "#805ad5")
    d = json.load(open(OUT / "eval_bestof3.json"))
    add("best-of-3 sampling (~matched FLOPs)", d["by_label"]["hard"], NH,
        "#805ad5")
    d = json.load(open(OUT / "eval_plan_baseline.json"))
    add("explicit plan in context (ceiling)", d["by_label"]["hard"], NH,
        "#1a202c")

    fig, ax = plt.subplots(figsize=(7.4, 4.8))
    style(ax)
    ys = range(len(rows))[::-1]
    for y, (label, h, n, color, note) in zip(ys, rows):
        lo, hi = wilson(round(h * n), n)
        ax.barh(y, h, color=color, height=0.62, alpha=0.88)
        ax.plot([lo, hi], [y, y], color="#333", lw=1.2)
        txt = f"{h:.2f}"
        if note:
            txt += f"  ({note})"
        ax.text(hi + 0.015, y, txt, va="center", fontsize=8)
    ax.set_yticks(list(ys))
    ax.set_yticklabels([r[0] for r in rows], fontsize=9)
    ax.set_xlim(0, 1.02)
    ax.set_xlabel("pass@1, MBPP hard bucket (plan-dependent items)")
    ax.set_title("Attribution ladder: what closes the plan gap "
                 "(bars: point estimate, whiskers: Wilson 95%)", fontsize=11)
    fig.tight_layout()
    fig.savefig(OUT / "fig_ladder.png", dpi=140, facecolor="white",
                bbox_inches="tight")
    print("wrote fig_ladder.png")


# ---------------- fig 3: cross-scale grid ----------------
def fig_scale():
    N2 = ROOT / "results-node2-final/results-12b"
    N1 = ROOT / "results-node-final/results-12b"
    panels = {
        ("MBPP", "E2B"): [
            ("base", *ev(OUT / "eval_code_code_s0_full.json", 0)),
            ("loop k=4", *ev(OUT / "eval_code_code_s0_full.json", 4)),
            ("adaptive k=4", *ev(OUT / "eval_code_e2b_adaptive.json", 4)),
            ("distill", *ev(OUT / "eval_code_distill_s1.json", 1)),
        ],
        ("MBPP", "12B"): [
            ("base", *ev(N1 / "eval_code_12b_trained.json", 0)),
            ("loop k=4 (α=.3)", *ev(N1 / "eval_code_12b_trained.json", 4)),
            ("adaptive k=4", *ev(N2 / "eval_code_12b_adaptive.json", 4)),
            ("distill", *ev(OUT / "eval_code_12b_distill.json", 1)),
        ],
        ("GSM8K", "E2B"): [
            ("base", *ev(OUT / "eval_uni.json", 0)),
            ("loop k=2", *ev(OUT / "eval_uni.json", 2)),
            ("adaptive k=2", *ev(OUT / "eval_gsm_e2b_adaptive.json", 2)),
            ("distill", *ev(OUT / "eval_gsm_distill_retry.json", 1)),
        ],
        ("GSM8K", "12B"): [
            ("base", *ev(N1 / "eval_12b_gsm_loop.json", 0)),
            ("loop k=2 (α=.3)", *ev(N1 / "eval_12b_gsm_loop.json", 2)),
            ("adaptive k=2", *ev(OUT / "eval_12b_gsm_adaptive.json", 2)),
            ("distill*", *ev(OUT / "eval_12b_gsm_distill.json", 1)),
        ],
    }
    fig, axes = plt.subplots(2, 2, figsize=(9.6, 6.6))
    for ax, ((task, scale), arms) in zip(axes.flat, panels.items()):
        if (task, scale) == ("GSM8K", "12B"):
            ax.annotate("*training collapse (0.00)", (3, 0.03), fontsize=7.5,
                        ha="center", color=RED)
        style(ax)
        x = range(len(arms))
        ax.bar([i - 0.19 for i in x], [a[1] for a in arms], width=0.36,
               color=GRAY, alpha=0.85, label="overall")
        ax.bar([i + 0.19 for i in x], [a[2] for a in arms], width=0.36,
               color=BLUE, alpha=0.85, label="hard")
        ax.axhline(arms[0][1], color=GRAY, lw=1, ls=":")
        ax.set_xticks(list(x))
        ax.set_xticklabels([a[0] for a in arms], fontsize=8)
        ax.set_title(f"{task} · {scale}", fontsize=10)
        ax.set_ylim(0, 1.0)
        if ax is axes.flat[0]:
            ax.legend(fontsize=8, frameon=False)
    fig.suptitle("Cross-scale attribution: constant-α destroys the 12B "
                 "substrate; state-dependent α restores it (MBPP) but not "
                 "everywhere", fontsize=11.5)
    fig.tight_layout(rect=(0, 0, 1, 0.96))
    fig.savefig(OUT / "fig_scale.png", dpi=140, facecolor="white",
                bbox_inches="tight")
    print("wrote fig_scale.png")


# ---------------- fig 4: transfer panel ----------------
def fig_transfer():
    def he_hard(path, k):
        d = json.load(open(OUT / path))
        return d["ks"][str(k)]["acc"], d["ks"][str(k)]["hard_acc"]

    groups = [
        ("HumanEval\n(MBPP-trained)", [
            ("base", *he_hard("eval_humaneval_trained.json", 0)),
            ("loop k=4", *he_hard("eval_humaneval_trained.json", 4)),
            ("distill", *he_hard("eval_humaneval_distill_transfer.json", 1)),
        ]),
        ("Rust / MultiPL-E\n(Python-trained)", [
            ("base", *ev(OUT / "eval_rust_py_transfer.json", 0)),
            ("loop k=4", *ev(OUT / "eval_rust_py_transfer.json", 4)),
        ]),
        ("LiveCodeBench (far)\nuntrained wins", [
            ("base", *he_hard("eval_lcb_untrained.json", 0)),
            ("untrained k=4", *he_hard("eval_lcb_untrained.json", 4)),
            ("loop k=4", *he_hard("eval_lcb_loop.json", 4)),
            ("distill", *he_hard("eval_lcb_distill.json", 1)),
        ]),
    ]
    fig, axes = plt.subplots(1, 3, figsize=(10.8, 3.8))
    for ax, (title, arms) in zip(axes, groups):
        style(ax)
        x = range(len(arms))
        ax.bar([i - 0.19 for i in x], [a[1] for a in arms], width=0.36,
               color=GRAY, alpha=0.85, label="overall")
        ax.bar([i + 0.19 for i in x],
               [a[2] if a[2] is not None else 0 for a in arms],
               width=0.36, color=BLUE, alpha=0.85, label="hard")
        ax.set_xticks(list(x))
        ax.set_xticklabels([a[0] for a in arms], fontsize=8.5)
        ax.set_title(title, fontsize=9.5)
        ax.set_ylim(0, 1.0)
    axes[0].legend(fontsize=8, frameon=False)
    fig.suptitle("Transfer: the implant moves with the substrate, "
                 "not the task", fontsize=11.5)
    fig.tight_layout(rect=(0, 0, 1, 0.94))
    fig.savefig(OUT / "fig_transfer.png", dpi=140, facecolor="white",
                bbox_inches="tight")
    print("wrote fig_transfer.png")


if __name__ == "__main__":
    fig_placement()
    fig_ladder()
    fig_scale()
    fig_transfer()