#!/usr/bin/env python3
"""Build the paper's two figures and their underlying data from
`results/reproduced_headline.json` (or any headline JSON of the same schema).

    python3 eval/figures.py [--headline results/published/reproduced_headline.json]

Writes results/figure_headline.{png,pdf}, results/figure_causal.{png,pdf} and
results/figure_headline_data.csv — the numbers the headline figure plots, so the
figure is checkable without rendering it.
"""
import argparse
import csv
import json
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt  # noqa: E402
import numpy as np  # noqa: E402

ROOT = Path(__file__).resolve().parent.parent
E1_ORDER = [("PNSR_K4_fin", "PNSR s1"), ("PNSR_K4_s2", "PNSR s2"),
            ("PNSR_K4_s3", "PNSR s3"), ("G_TAU512", "τ=512"),
            ("G_CONVEX_FREE", "gated free"), ("G_CONVEX_RETAIN", "gated retain"),
            ("RMT_s1", "RMT")]
HANDLES = [("J_self_lesion_SELF_REF", "exact store: J_self lesion", "#1b6f4a", "exact"),
           ("J_obs_lesion_pointer_families", "exact store: J_obs lesion", "#3aa06b", "exact"),
           ("delta_cross_lifetime_swap", "learned state: cross-lifetime swap", "#7b2d5e", "learned"),
           ("delta_reset64", "learned state: reset every 64", "#b8547f", "learned"),
           ("delta_state_zero", "learned state: zero + freeze writes", "#d98cae", "learned")]


def load(path):
    return json.loads(Path(path).read_text())


def headline(H, out):
    e1m = H["experiment1_generic_recurrent_state"]["models"]
    e2 = H["experiment2_protected_bindings"]
    e1 = [(lab, e1m[k]["learned_state_causal"]["delta_cross_lifetime_swap"])
          for k, lab in E1_ORDER if k in e1m]
    unb = [(f"unbound s{s}", e2["capacity_matched_unbound"][f"seed{s}"]["delta_payload_swap"])
           for s in (1, 2, 3) if f"seed{s}" in e2.get("capacity_matched_unbound", {})]
    bnd = [(f"bind s{s}", e2["bind"][f"seed{s}"]["delta_payload_swap"])
           for s in (1, 2, 3) if f"seed{s}" in e2.get("bind", {})]

    with open(out / "figure_headline_data.csv", "w", newline="") as f:
        w = csv.writer(f)
        w.writerow(["group", "label", "delta_cross_lifetime_swap"])
        for g, items in (("experiment1_unstructured_learned_state", e1),
                         ("experiment2_capacity_matched_unstructured_control", unb),
                         ("experiment2_binding_local_writes", bnd)):
            for lab, v in items:
                w.writerow([g, lab, v])

    fig, ax = plt.subplots(figsize=(12.5, 5.6), constrained_layout=True)
    groups = [("Experiment 1\nunstructured learned state\n(4 mechanisms, 3 seeds)", e1, "#8a8f98"),
              ("Experiment 2 control\ncapacity-matched,\nunstructured writes", unb, "#c8642a"),
              ("Experiment 2\nbinding-local writes", bnd, "#1b6f4a")]
    x, ticks, labels, centres = 0, [], [], []
    for gname, items, col in groups:
        xs = []
        for n, v in items:
            ax.bar(x, v, 0.72, color=col, edgecolor="black", linewidth=0.5)
            ticks.append(x)
            labels.append(n)
            xs.append(x)
            x += 1
        if xs:
            centres.append((float(np.mean(xs)), gname))
        x += 1.1
    ax.axhline(0, color="black", lw=0.9)
    ax.axhline(0.10, color="#1b6f4a", ls="--", lw=1.1,
               label="preregistered causal threshold (0.10)")
    ax.set_xticks(ticks)
    ax.set_xticklabels(labels, fontsize=9, rotation=35, ha="right")
    ax.set_ylabel("causal effect of the learned state\n(warm − cross-lifetime swap)", fontsize=11)
    ax.set_title("A learned recurrent state becomes causally load-bearing only when "
                 "entity bindings are structurally protected", fontsize=13, weight="bold")
    ax.legend(fontsize=9, loc="upper left")
    ax.grid(axis="y", alpha=0.25)
    ax.set_ylim(-0.02, 0.20)
    for cx, gname in centres:
        ax.annotate(gname, xy=(cx, -0.0165), xycoords=("data", "data"), ha="center",
                    fontsize=9.5, weight="bold", annotation_clip=False)
    fig.savefig(out / "figure_headline.png", dpi=200)
    fig.savefig(out / "figure_headline.pdf")
    plt.close(fig)
    print("  figure_headline: E1", [round(v, 4) for _, v in e1])
    print("                   unbound", [round(v, 4) for _, v in unb],
          "bind", [round(v, 4) for _, v in bnd])


def causal(H, out):
    m = H["experiment1_generic_recurrent_state"]["models"]
    runs = [(k, lab) for k, lab in E1_ORDER if k in m]
    fig, (a, b) = plt.subplots(1, 2, figsize=(14, 5.4), constrained_layout=True,
                               gridspec_kw={"width_ratios": [2.05, 1]})
    n, w = len(HANDLES), 0.16
    for i, (key, lab, col, _) in enumerate(HANDLES):
        xs = np.arange(len(runs)) + (i - (n - 1) / 2) * w
        ys = [(m[k].get("exact_store_lesions", {}) | m[k].get("learned_state_causal", {})
               ).get(key) or 0.0 for k, _ in runs]
        a.bar(xs, ys, w, label=lab, color=col, edgecolor="black", linewidth=0.4)
    a.axhline(0, color="black", lw=0.9)
    a.set_xticks(np.arange(len(runs)))
    a.set_xticklabels([lab for _, lab in runs], fontsize=9, rotation=20, ha="right")
    a.set_ylabel("causal effect size (warm − intervened accuracy)", fontsize=11)
    a.set_title("Exact stores are causally necessary; the learned state is not",
                fontsize=12, weight="bold")
    a.legend(fontsize=8.5, ncol=2)
    a.grid(axis="y", alpha=0.25)

    cos = [m[k].get("state_statistics", {}).get("mean_cross_lifetime_cosine") for k, _ in runs]
    nrm = [m[k].get("state_statistics", {}).get("max_slot_norm") for k, _ in runs]
    b.scatter(cos, nrm, s=90, c="#7b2d5e", edgecolor="black", zorder=3)
    for (k, lab), c, nn in zip(runs, cos, nrm):
        if c is not None and nn is not None:
            b.annotate(lab.replace("\n", " "), (c, nn), fontsize=8,
                       xytext=(4, 4), textcoords="offset points")
    b.set_yscale("log")
    b.set_xlabel("mean cross-lifetime cosine\n(is the state lifetime-specific?)", fontsize=10)
    b.set_ylabel("max slot norm (log)", fontsize=10)
    b.set_title("States differ enormously; none is causally used", fontsize=12, weight="bold")
    b.grid(alpha=0.25)
    fig.savefig(out / "figure_causal.png", dpi=200)
    fig.savefig(out / "figure_causal.pdf")
    plt.close(fig)
    print("  figure_causal: cosine", [c for c in cos], "norms", [nn for nn in nrm])


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--headline", default=str(ROOT / "results" / "published" /
                                              "reproduced_headline.json"))
    ap.add_argument("--outdir", default=str(ROOT / "results"))
    args = ap.parse_args()
    out = Path(args.outdir)
    out.mkdir(parents=True, exist_ok=True)
    H = load(args.headline)
    headline(H, out)
    causal(H, out)
    print("->", out)


if __name__ == "__main__":
    main()