"""Publication figures from analysis/tables and data/validated (vector PDF, matplotlib). Palette = dataviz reference slots 1-5 (validated).
Usage: .venv/bin/python analysis/figures/make_figures.py --tag pilot"""
from __future__ import annotations
import argparse, json, os
import numpy as np, pandas as pd
import matplotlib; matplotlib.use("Agg")
import matplotlib.pyplot as plt
REPO = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))
TAB = os.path.join(REPO, "analysis", "tables"); FIG = os.path.join(REPO, "analysis", "figures", "out"); VAL = os.path.join(REPO, "data", "validated")
ARMS = ["A0", "A1", "A2", "A3", "A4"]; COL = {"A0": "#2a78d6", "A1": "#eb6834", "A2": "#1baf7a", "A3": "#eda100", "A4": "#e87ba4"}
LABEL = {"A0": "A0 restricted", "A1": "A1 shell", "A2": "A2 macOS", "A3": "A3 computer-use", "A4": "A4 self-repair"}
plt.rcParams.update({"font.family": "serif", "font.size": 9, "axes.spines.top": False, "axes.spines.right": False, "axes.grid": True, "grid.color": "#e6e6e6", "grid.linewidth": 0.5, "axes.edgecolor": "#888", "legend.frameon": False, "figure.dpi": 150})

def save(fig, name):
    os.makedirs(FIG, exist_ok=True); fig.savefig(os.path.join(FIG, name + ".pdf"), bbox_inches="tight"); fig.savefig(os.path.join(FIG, name + ".png"), bbox_inches="tight", dpi=200); plt.close(fig)

def fig_completion_by_arm(t1, t2, tag):
    fig, ax = plt.subplots(figsize=(5.2, 3.0))
    x = np.arange(len(t1)); ax.bar(x, t1["acceptance_taskw"], color=[COL[a] for a in t1.arm], width=0.55)
    ax.bar(x + 0.0, t1["det_completion"], facecolor="none", edgecolor="#444", width=0.55, linewidth=0.8, linestyle="--", label="deterministic completion")
    for i, v in enumerate(t1["acceptance_taskw"]): ax.text(i, v + 0.02, f"{v:.0%}", ha="center", fontsize=8)
    ax.set_xticks(x); ax.set_xticklabels([LABEL[a] for a in t1.arm], rotation=0, fontsize=8); ax.set_ylim(0, 1.05); ax.set_ylabel("scenario-weighted acceptance rate")
    ax.legend(loc="upper left", fontsize=8); save(fig, f"fig03_acceptance_by_arm_{tag}")

def fig_contrasts(t2, tag):
    d = t2[t2.outcome == "S"]; fig, ax = plt.subplots(figsize=(5.2, 2.6)); y = np.arange(len(d))
    ax.errorbar(d["estimate"], y, xerr=[d["estimate"] - d["ci_lo"], d["ci_hi"] - d["estimate"]], fmt="o", color="#2a78d6", ecolor="#2a78d6", capsize=3, markersize=5)
    ax.axvline(0, color="#888", linewidth=0.8); ax.set_yticks(y); ax.set_yticklabels([f"{c} ({k})" for c, k in zip(d.contrast, d.kind)]); ax.set_xlabel("difference in acceptance probability (scenario-cluster bootstrap 95% CI)")
    for yi, (e, p) in enumerate(zip(d["estimate"], d["perm_p"])): ax.text(e, yi + 0.28, f"{e:+.2f}, p={p:.2f}", fontsize=7, ha="center")
    save(fig, f"fig07_tool_effects_{tag}")

def fig_occupation(t3, tag):
    piv = t3.pivot(index="occupation", columns="arm", values="acceptance").reindex(columns=[a for a in ARMS if a in t3.arm.unique()])
    fig, ax = plt.subplots(figsize=(5.6, 3.6)); im = ax.imshow(piv.values, cmap=matplotlib.colors.LinearSegmentedColormap.from_list("b", ["#f5f8fd", "#cde2fb", "#6da7ec", "#256abf", "#0d366b"]), vmin=0, vmax=1, aspect="auto")
    ax.set_xticks(range(piv.shape[1])); ax.set_xticklabels(piv.columns); ax.set_yticks(range(piv.shape[0])); ax.set_yticklabels(piv.index, fontsize=8); ax.grid(False)
    for i in range(piv.shape[0]):
        for j in range(piv.shape[1]):
            v = piv.values[i, j]
            if v == v: ax.text(j, i, f"{v:.0%}", ha="center", va="center", fontsize=7, color="white" if v > 0.55 else "#0b0b0b")
    fig.colorbar(im, ax=ax, fraction=0.03, label="acceptance rate"); save(fig, f"fig14_occupation_heterogeneity_{tag}")

def fig_cost_quality(df, tag):
    fig, ax = plt.subplots(figsize=(5.2, 3.2))
    for a in ARMS:
        d = df[df.arm == a]
        if len(d) == 0: continue
        ax.scatter(d["C"], d["quality"], s=22, color=COL[a], alpha=0.75, edgecolor="white", linewidth=0.5, label=LABEL[a], marker="o" if a != "A4" else "s")
        ax.scatter(d["C"].mean(), d["quality"].mean(), s=90, color=COL[a], edgecolor="#222", linewidth=1.0, marker="D")
    ax.axhline(70, color="#888", linewidth=0.8, linestyle=":"); ax.text(ax.get_xlim()[1], 70.5, "q = 70", fontsize=7, ha="right")
    ax.set_xlabel("marginal API cost per run (USD, incl. recovery)"); ax.set_ylabel("blind quality score (0–100)"); ax.legend(fontsize=7, ncol=2); save(fig, f"fig05_cost_quality_{tag}")

def fig_alf(t5, tag):
    d = t5[t5.weights == "equal"]; occs = sorted(d.occupation.unique()); fig, axes = plt.subplots(2, int(np.ceil(len(occs) / 2)), figsize=(7.5, 4.2), sharex=True, sharey=True); axes = np.array(axes).ravel()
    for ax, occ in zip(axes, occs):
        for a in ARMS:
            dd = d[(d.occupation == occ) & (d.arm == a)]
            if len(dd): ax.plot(dd["B"], dd["ALF"], color=COL[a], linewidth=1.6, marker="o", markersize=3, label=LABEL[a])
        ax.set_title(occ, fontsize=8); ax.set_ylim(0, 1.02); ax.set_xscale("log"); ax.set_xticks([0.25, 0.5, 1, 2, 4]); ax.set_xticklabels(["0.25", "0.5", "1", "2", "4"], fontsize=7)
    for ax in axes[len(occs):]: ax.axis("off")
    axes[0].legend(fontsize=6, loc="upper left"); fig.supxlabel("budget cap B (USD per run)", fontsize=8); fig.supylabel("ALF$_o$(B, H=0, q=70)", fontsize=8); save(fig, f"fig04_alf_curves_{tag}")

def fig_time(df, tag):
    fig, ax = plt.subplots(figsize=(5.2, 3.0))
    for a in ARMS:
        d = df[df.arm == a]
        if len(d) == 0: continue
        t = np.sort(d["T"].values); ax.step(t, np.arange(1, len(t) + 1) / len(t), where="post", color=COL[a], label=LABEL[a], linewidth=1.6)
    ax.axvline(1800, color="#888", linestyle=":", linewidth=0.8); ax.set_xlabel("wall time per run (s)"); ax.set_ylabel("empirical CDF"); ax.legend(fontsize=7); save(fig, f"fig11_time_cdf_{tag}")

def fig_dependency(df, tag):
    fig, ax = plt.subplots(figsize=(5.2, 3.0))
    for a in ARMS:
        d = df[df.arm == a]
        if len(d) == 0: continue
        g = d.groupby("tool_calls")["S"].mean()
        ax.scatter(d["tool_calls"], d["S"] + np.random.default_rng(1).normal(0, 0.015, len(d)), s=14, color=COL[a], alpha=0.6, label=LABEL[a])
    ax.set_xlabel("tool calls in the run"); ax.set_ylabel("accepted (jittered)"); ax.legend(fontsize=7); save(fig, f"fig06_dependency_{tag}")

def fig_failures(df, tag):
    cats = ["accepted", "det. pass, quality < q", "det. fail, no violation", "critical violation", "timeout / no result"]
    fig, ax = plt.subplots(figsize=(5.2, 3.0)); bottom = np.zeros(len(ARMS)); cols = ["#2a78d6", "#86b6ef", "#c3c2b7", "#d03b3b", "#52514e"]
    for ci, c in enumerate(cats):
        vals = []
        for a in ARMS:
            d = df[df.arm == a]
            if len(d) == 0: vals.append(0); continue
            if c == "accepted": v = (d["S"] == 1).mean()
            elif c == "det. pass, quality < q": v = ((d["D"] == 1) & (d["E"] == 0) & (d["S"] == 0)).mean()
            elif c == "det. fail, no violation": v = ((d["D"] == 0) & (d["E"] == 0) & (~d["censored"])).mean()
            elif c == "critical violation": v = (d["E"] == 1).mean()
            else: v = ((d["censored"]) & (d["D"] == 0) & (d["E"] == 0)).mean()
            vals.append(v)
        ax.bar(range(len(ARMS)), vals, bottom=bottom, color=cols[ci], label=c, width=0.6, edgecolor="white", linewidth=1); bottom += np.array(vals)
    ax.set_xticks(range(len(ARMS))); ax.set_xticklabels(ARMS); ax.set_ylabel("share of runs"); ax.legend(fontsize=7, loc="upper left", bbox_to_anchor=(1, 1)); save(fig, f"fig08_failure_severity_{tag}")

def fig_judges(jdg, tag):
    piv = jdg.pivot_table(index="run_id", columns="judge_idx", values="score")
    if piv.shape[1] < 2: return
    fig, axes = plt.subplots(1, piv.shape[1] - 1, figsize=(2.6 * (piv.shape[1] - 1), 2.6))
    axes = np.atleast_1d(axes)
    for k, ax in enumerate(axes):
        ax.scatter(piv[0], piv[k + 1], s=14, color="#2a78d6", alpha=0.7); ax.plot([0, 100], [0, 100], color="#888", linewidth=0.8); ax.set_xlabel("judge 1 (Opus)"); ax.set_ylabel(f"judge {k + 2} (Sonnet)"); ax.set_xlim(0, 100); ax.set_ylim(0, 100)
    save(fig, f"fig12_judge_agreement_{tag}")

def fig_recovery(df, tag):
    a4 = df[df.arm == "A4"]
    if len(a4) == 0: return
    f = a4[a4.selfcheck_pass_first == 0]; fig, ax = plt.subplots(figsize=(3.6, 2.6))
    vals = [len(a4), (a4.selfcheck_pass_first == 1).sum(), len(f), (f.selfcheck_pass_final == 1).sum(), int(f["S"].sum())]
    ax.bar(range(5), vals, color=["#2a78d6", "#1baf7a", "#eda100", "#1baf7a", "#2a78d6"], width=0.6)
    ax.set_xticks(range(5)); ax.set_xticklabels(["A4 runs", "pass 1st\nself-check", "fail 1st", "recovered\n(pass final)", "accepted\nafter recovery"], fontsize=7)
    for i, v in enumerate(vals): ax.text(i, v + 0.2, str(v), ha="center", fontsize=8)
    ax.set_ylabel("runs"); save(fig, f"fig09_recovery_{tag}")

def main():
    ap = argparse.ArgumentParser(); ap.add_argument("--tag", default="pilot"); a = ap.parse_args(); tag = a.tag
    import sys; sys.path.insert(0, os.path.join(REPO, "analysis", "econometrics")); from estimate import load
    df = load(tag, 70); t1 = pd.read_csv(os.path.join(TAB, f"table1_outcomes_by_arm_{tag}.csv")); t2 = pd.read_csv(os.path.join(TAB, f"table2_contrasts_{tag}.csv"))
    t3 = pd.read_csv(os.path.join(TAB, f"table3_occupation_arm_{tag}.csv")); t5 = pd.read_csv(os.path.join(TAB, f"table5_alf_{tag}.csv")); jdg = pd.read_parquet(os.path.join(VAL, f"judgments_{tag}.parquet"))
    fig_completion_by_arm(t1, t2, tag); fig_contrasts(t2, tag); fig_occupation(t3, tag); fig_cost_quality(df, tag); fig_alf(t5, tag); fig_time(df, tag); fig_dependency(df, tag); fig_failures(df, tag); fig_judges(jdg, tag); fig_recovery(df, tag)
    print("figures written to", FIG)

if __name__ == "__main__":
    main()