#!/usr/bin/env python3
"""Regenerate intervention figures for the README from saved JSON."""
from __future__ import annotations

import json
from pathlib import Path

import matplotlib

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

ROOT = Path(__file__).resolve().parents[1]
IN = ROOT / "figs" / "interventions"
OUT = IN

COLORS = {"CSO": "#c0392b", "BT W=5": "#e67e22", "BT W=2": "#2980b9"}
ORDER = ["CSO", "BT W=5", "BT W=2"]


def load_methods():
    out = {}
    for name, fname in [
        ("CSO", "L20_CSO.json"),
        ("BT W=5", "L20_BT_W5.json"),
        ("BT W=2", "L20_BT_W2.json"),
    ]:
        out[name] = json.load(open(IN / fname))
    return out


def plot_all_methods(methods):
    fig, axes = plt.subplots(2, 2, figsize=(12.5, 8.5))
    fig.suptitle(
        "L20 interventions: CSO vs BT W=5 vs BT W=2 (128 graphs)",
        fontsize=13,
        fontweight="bold",
    )

    # A: pinlast-k
    ax = axes[0, 0]
    for name in ORDER:
        pk = methods[name]["pinlastk"]
        ks = [0] + list(pk["ks"])
        ys = [pk["clean"]] + list(pk["rand_mean"])
        ax.plot(ks, ys, "o-", color=COLORS[name], label=name, lw=2, ms=5)
    ax.set_xlabel("k earlier latents corrupted (last pinned)")
    ax.set_ylabel("leaf accuracy")
    ax.set_ylim(0, 1.05)
    ax.set_title("A. Pin-last: corrupt k earlier thoughts")
    ax.grid(True, alpha=0.3)
    ax.legend(frameon=False)

    # B: aggregate bars
    ax = axes[0, 1]
    keys = [
        ("clean", "clean"),
        ("earlier L1..L19 -> noise (last pinned)", "earlier→noise\n(last pin)"),
        ("earlier L1..L19 -> other-graph (last pinned)", "earlier→donor\n(last pin)"),
        ("LAST thought L20 -> other-graph", "last→donor"),
        ("PROPAGATED: corrupt L10, recompute rest", "prop L10"),
        ("PROPAGATED: corrupt L19, recompute rest", "prop L19"),
    ]
    x = np.arange(len(keys))
    width = 0.25
    for i, name in enumerate(ORDER):
        rows = {r["label"]: r["acc"] for r in methods[name]["aggregate"]["rows"]}
        ys = [rows[k] for k, _ in keys]
        ax.bar(x + (i - 1) * width, ys, width, color=COLORS[name], label=name)
    ax.set_xticks(x)
    ax.set_xticklabels([lab for _, lab in keys], fontsize=8)
    ax.set_ylabel("leaf accuracy")
    ax.set_ylim(0, 1.05)
    ax.set_title("B. Aggregate interventions")
    ax.legend(frameon=False, fontsize=8)
    ax.grid(True, axis="y", alpha=0.3)

    # C: per-slot pin-last
    ax = axes[1, 0]
    for name in ORDER:
        ys = methods[name]["perslot"]["pin_accs"]
        xs = np.arange(1, len(ys) + 1)
        ax.plot(xs, ys, "o-", color=COLORS[name], label=name, lw=2, ms=4)
    ax.axvline(20, color="gray", ls="--", alpha=0.5)
    ax.set_xlabel("corrupted latent slot (last pinned unless slot=20)")
    ax.set_ylabel("leaf accuracy")
    ax.set_ylim(0, 1.05)
    ax.set_title("C. Per-slot pin-last")
    ax.grid(True, alpha=0.3)
    ax.legend(frameon=False)

    # D: per-slot propagate
    ax = axes[1, 1]
    for name in ORDER:
        ys = methods[name]["perslot"]["prop_accs"]
        xs = np.arange(1, len(ys) + 1)
        ax.plot(xs, ys, "o-", color=COLORS[name], label=name, lw=2, ms=4)
    ax.set_xlabel("corrupted latent slot (then recompute rest)")
    ax.set_ylabel("leaf accuracy")
    ax.set_ylim(0, 1.05)
    ax.set_title("D. Per-slot propagate")
    ax.grid(True, alpha=0.3)
    ax.legend(frameon=False)

    fig.tight_layout(rect=[0, 0, 1, 0.96])
    path = OUT / "L20_all_methods.png"
    fig.savefig(path, dpi=160)
    plt.close(fig)
    print(f"saved {path}")


def plot_pinlast_k(methods):
    fig, ax = plt.subplots(figsize=(7.2, 4.2))
    for name in ORDER:
        pk = methods[name]["pinlastk"]
        ks = [0] + list(pk["ks"])
        ys = [pk["clean"]] + list(pk["rand_mean"])
        ax.plot(ks, ys, "o-", color=COLORS[name], label=name, lw=2.2, ms=6)
    ax.set_title("L20: pin-last intervention")
    ax.set_xlabel("k earlier latents corrupted (last pinned)")
    ax.set_ylabel("leaf accuracy")
    ax.set_ylim(0, 1.05)
    ax.grid(True, alpha=0.3)
    ax.legend(frameon=False)
    fig.tight_layout()
    path = OUT / "pinlast_k_L20.png"
    fig.savefig(path, dpi=160)
    plt.close(fig)
    print(f"saved {path}")


def plot_pinlast_table(methods):
    ks = methods["BT W=2"]["pinlastk"]["ks"]
    rows = ["clean"] + [str(k) for k in ks] + ["all 19"]
    col_labels = ["k"] + ORDER
    cell = []
    for name in ORDER:
        pk = methods[name]["pinlastk"]
        vals = [pk["clean"]] + list(pk["rand_mean"]) + [pk["all"]]
        cell.append([f"{v:.3f}" for v in vals])
    data = [[rows[i]] + [cell[j][i] for j in range(3)] for i in range(len(rows))]

    fig, ax = plt.subplots(figsize=(7.5, 4.8))
    ax.axis("off")
    ax.set_title(
        "L20 pin-last: k earlier thoughts corrupted (last pinned)",
        fontsize=11,
        pad=12,
    )
    table = ax.table(
        cellText=data,
        colLabels=col_labels,
        loc="center",
        cellLoc="center",
    )
    table.auto_set_font_size(False)
    table.set_fontsize(9)
    table.scale(1.15, 1.35)
    for (r, c), cell_obj in table.get_celld().items():
        if r == 0:
            cell_obj.set_facecolor("#ecf0f1")
            cell_obj.set_text_props(fontweight="bold")
        if c == 0:
            cell_obj.set_text_props(fontweight="bold")
    ckpts = ", ".join(
        f"{n}={Path(methods[n]['ckpt']).name.split('_')[-1]}" for n in ORDER
    )
    fig.text(0.5, 0.02, f"ckpts: {ckpts}", ha="center", fontsize=8, color="#555")
    fig.tight_layout(rect=[0, 0.05, 1, 1])
    path = OUT / "pinlast_k_L20_table.png"
    fig.savefig(path, dpi=160)
    plt.close(fig)
    print(f"saved {path}")


def main():
    methods = load_methods()
    plot_all_methods(methods)
    plot_pinlast_k(methods)
    plot_pinlast_table(methods)


if __name__ == "__main__":
    main()