"""Quantitative figures (vector PDF) generated from the frozen analysis dataset. Every figure shows n and uncertainty where relevant."""
import os, sys, json
sys.path.insert(0, os.path.dirname(__file__))
from common import *
import matplotlib; matplotlib.use("Agg")
import matplotlib.pyplot as plt
FIG = os.path.join(REPO, "results", "figures"); os.makedirs(FIG, exist_ok=True)
COL = {"A": "#7f7f7f", "B": "#4c72b0", "C": "#dd8452", "D": "#55a868"}
plt.rcParams.update({"font.size": 9, "axes.spines.top": False, "axes.spines.right": False, "pdf.fonttype": 42})

def fig_primary(split):
    at = pd.read_csv(os.path.join(EST, f"{split}_arm_table.csv")); ct = pd.read_csv(os.path.join(EST, f"{split}_contrasts.csv")); ct = ct[ct["endpoint"] == "Y"]
    df = load_runs(split)
    fig, ax = plt.subplots(1, 2, figsize=(7.4, 3.1), gridspec_kw={"width_ratios": [1, 1.5]})
    # left: task-weighted rates with cluster-bootstrap CI (per arm)
    for i, a in enumerate(ARMS):
        d = df[df["treatment"] == a].groupby("task_id")["Y"].mean()
        est, lo, hi = cluster_bootstrap(d - 0, 3000)
        ax[0].bar(i, est, color=COL[a], width=0.65); ax[0].errorbar(i, est, yerr=[[est - lo], [hi - est]], color="k", capsize=3, lw=1)
        ax[0].text(i, min(hi + 0.03, 1.02), f"{len(d)} tasks\n{len(df[df['treatment']==a])} runs", ha="center", va="bottom", fontsize=6.5)
    ax[0].set_xticks(range(4)); ax[0].set_xticklabels([f"{a}" for a in ARMS]); ax[0].set_xlabel("Arm"); ax[0].set_ylim(0, 1.2); ax[0].set_ylabel("Verified completion $\\hat p_a$ (task-weighted)")
    ax[0].set_title("(a) Primary endpoint by arm", loc="left")
    # right: paired contrasts
    cts = ct.reset_index(drop=True)
    for i, r in cts.iterrows():
        ax[1].errorbar(r["estimate"], i, xerr=[[r["estimate"] - r["ci_low"]], [r["ci_high"] - r["estimate"]]], fmt="o", color="k" if r["role"] == "primary" else "#555", capsize=3)
        ax[1].text(1.02, i, f"p={r['perm_p']:.2f}\n+/=/−: {int(r['n_tasks_favoring_first'])}/{int(r['n_tasks_tied'])}/{int(r['n_tasks_favoring_second'])}", fontsize=5.8, va="center", ha="left", transform=ax[1].get_yaxis_transform(), color="#444")
    ax[1].axvline(0, color="#999", lw=0.8); ax[1].set_yticks(range(len(cts))); ax[1].set_yticklabels([f"{c} ({r})" for c, r in zip(cts["contrast"], cts["role"])]); ax[1].invert_yaxis()
    ax[1].set_xlabel("Paired task-level difference (95% task-cluster bootstrap CI)"); ax[1].set_title("(b) Prespecified contrasts", loc="left")
    fig.tight_layout(rect=(0, 0, 0.93, 1)); fig.savefig(os.path.join(FIG, f"fig_primary_{split}.pdf")); plt.close(fig)

def fig_quality(split):
    df = load_runs(split)
    comps = ["q_correctness", "q_integrity", "q_reproducibility", "q_reporting", "q_baseline"]; labels = ["correctness", "integrity", "reproducibility", "reporting", "baseline"]
    fig, ax = plt.subplots(1, 2, figsize=(7.2, 2.8))
    w = 0.2
    for i, a in enumerate(ARMS):
        d = df[df["treatment"] == a]; means = [d.groupby("task_id")[c].mean().mean() for c in comps]
        ax[0].bar(np.arange(len(comps)) + (i - 1.5) * w, means, width=w, color=COL[a], label=f"Arm {a} (n={len(d)})")
    ax[0].set_xticks(range(len(comps))); ax[0].set_xticklabels(labels, rotation=20); ax[0].set_ylim(0, 1.05); ax[0].set_ylabel("Mean component score (task-weighted)"); ax[0].legend(fontsize=6, frameon=False); ax[0].set_title("(a) Quality components", loc="left")
    # violations by type and arm
    cr = pd.read_csv(os.path.join(DATA, "criteria.csv")); cr = cr[cr["kind"] == "violation"].merge(df[["run_id", "treatment"]], on="run_id")
    if len(cr):
        tab = cr.groupby(["passed_or_type", "treatment"])["run_id"].nunique().unstack(fill_value=0).reindex(columns=ARMS, fill_value=0)
        n_runs = df.groupby("treatment").size().reindex(ARMS)
        rate = tab / n_runs
        rate.plot.bar(ax=ax[1], color=[COL[a] for a in ARMS], width=0.8); ax[1].set_ylabel("Share of runs with ≥1 violation of type"); ax[1].set_xlabel(""); ax[1].legend(fontsize=6, frameon=False, title=None)
        ax[1].set_xticklabels(ax[1].get_xticklabels(), rotation=20)
    ax[1].set_title("(b) Critical integrity violations by type", loc="left")
    fig.tight_layout(); fig.savefig(os.path.join(FIG, f"fig_quality_{split}.pdf")); plt.close(fig)

def fig_cost(split):
    df = load_runs(split)
    fig, ax = plt.subplots(1, 2, figsize=(7.2, 2.8))
    for a in ARMS:
        d = df[df["treatment"] == a]
        ax[0].scatter(d["spent_usd"], d["spent_wall_s"] / 60, s=14, color=COL[a], alpha=0.7, label=f"Arm {a}", marker="o" if True else "x")
        ax[1].scatter(d["spent_usd"].mean(), task_weighted_rate(d, a), s=60, color=COL[a], edgecolor="k", zorder=3)
        ax[1].annotate(f"{a}", (d["spent_usd"].mean(), task_weighted_rate(d, a)), xytext=(5, 4), textcoords="offset points", fontsize=8)
        ax[1].errorbar(d["spent_usd"].mean(), task_weighted_rate(d, a), xerr=d["spent_usd"].std() / np.sqrt(max(1, len(d))), color=COL[a], capsize=2, lw=0.8)
    ax[0].set_xlabel("Inference cost per run (USD, client estimate)"); ax[0].set_ylabel("Wall time per run (min)"); ax[0].legend(fontsize=6, frameon=False); ax[0].set_title("(a) Cost and time of every run", loc="left")
    ax[1].set_xlabel("Mean cost per run (USD) ± s.e."); ax[1].set_ylabel("Intervention-free verified completion"); ax[1].set_ylim(0, 1.05); ax[1].set_title("(b) Quality–cost frontier", loc="left")
    fig.tight_layout(); fig.savefig(os.path.join(FIG, f"fig_cost_{split}.pdf")); plt.close(fig)

def fig_family(split):
    fam = pd.read_csv(os.path.join(EST, f"{split}_family_table.csv"))
    fams = sorted(fam["family"].unique()); fig, ax = plt.subplots(figsize=(7.2, 2.6)); w = 0.2
    for i, a in enumerate(ARMS):
        d = fam[fam["treatment"] == a].set_index("family").reindex(fams)
        ax.bar(np.arange(len(fams)) + (i - 1.5) * w, d["Y"], width=w, color=COL[a], label=f"Arm {a}")
        for j, f in enumerate(fams):
            if i == 0: ax.text(j, 1.02, f"{int(d.loc[f, 'n_tasks'])} tasks", ha="center", fontsize=7)
    ax.set_xticks(range(len(fams))); ax.set_xticklabels(fams); ax.set_ylim(0, 1.12); ax.set_ylabel("Completion (task-weighted)"); ax.legend(fontsize=7, frameon=False, ncol=4); ax.set_title("Task-family heterogeneity of the primary endpoint", loc="left")
    fig.tight_layout(); fig.savefig(os.path.join(FIG, f"fig_family_{split}.pdf")); plt.close(fig)

def fig_errors(split):
    df = load_runs(split); cr = pd.read_csv(os.path.join(DATA, "criteria.csv")).merge(df[["run_id", "treatment", "task_id"]], on="run_id")
    fails = cr[(cr["kind"] == "criterion") & (cr["passed_or_type"].astype(str) == "0")]
    if fails.empty: return
    fails["error"] = fails["task_id"] + ":" + fails["id"]
    top = fails.groupby("error")["run_id"].nunique().sort_values(ascending=False).head(25)
    tab = fails[fails["error"].isin(top.index)].groupby(["error", "treatment"])["run_id"].nunique().unstack(fill_value=0).reindex(columns=ARMS, fill_value=0).loc[top.index]
    fig, ax = plt.subplots(figsize=(7.2, 0.22 * len(tab) + 1.2)); tab.plot.barh(ax=ax, stacked=True, color=[COL[a] for a in ARMS], width=0.8)
    ax.invert_yaxis(); ax.set_xlabel("Number of runs failing the criterion (denominator: runs of that task per arm)"); ax.legend(fontsize=7, frameon=False); ax.set_title("Most frequent failed scientific criteria (task:criterion)", loc="left")
    fig.tight_layout(); fig.savefig(os.path.join(FIG, f"fig_errors_{split}.pdf")); plt.close(fig)

if __name__ == "__main__":
    split = sys.argv[1] if len(sys.argv) > 1 else "confirmatory"
    for f in (fig_primary, fig_quality, fig_cost, fig_family, fig_errors):
        try: f(split); print("ok", f.__name__)
        except Exception as e: print("skip", f.__name__, e)