#!/usr/bin/env python3
"""Preregistered analysis of experiments/outputs/<protocol>/{runs,judgments}.jsonl → analysis tables (CSV/LaTeX) + figures. No data = no results."""
import json, sys, itertools
from pathlib import Path
import numpy as np, pandas as pd
from scipy import stats
sys.path.insert(0, "src")
from jof250 import config, viz
from jof250.provenance import stamp
PROTO = sys.argv[1] if len(sys.argv) > 1 else "heldout_v1"
R = config.ROOT; OUT = R / "experiments" / "outputs" / PROTO; AN = R / "experiments" / "analysis"; AN.mkdir(exist_ok=True); TAB = R / "paper" / "tables"; TAB.mkdir(exist_ok=True)
runs = [json.loads(l) for l in (OUT / "runs.jsonl").read_text().splitlines() if l.strip()]; judg = [json.loads(l) for l in (OUT / "judgments.jsonl").read_text().splitlines() if l.strip()] if (OUT / "judgments.jsonl").exists() else []
DIMS = ["clarity", "argument_logic", "contribution_positioning", "methodological_rigor", "fidelity", "usefulness"]
rows = []
for j in judg:
    if not j.get("ok"): continue
    for cond, sc in j["scores"].items():
        rows.append({"task": j["task"], "doc": j["doc"], "rep": j["rep"], "pass": j["pass"], "condition": cond, **{d: sc.get(d) for d in DIMS}, "n_invented_facts": sc.get("n_invented_facts")})
J = pd.DataFrame(rows); Rn = pd.DataFrame([{**{k: r[k] for k in ("task", "doc", "domain", "condition", "rep", "ok")}, **({k: r["metrics"][k] for k in ("n_words", "preserved_rate", "new_numbers", "n_unverifiable_citations", "n_unsupported_claims", "max_overlap_share", "todo_placeholders")} if r.get("ok") else {}), "cost_usd": r.get("cost_usd"), "latency_s": r.get("latency_s")} for r in runs])
res = {"protocol": PROTO, "generated": stamp(), "n_runs": len(Rn), "n_runs_ok": int(Rn.ok.sum()), "n_judgments": len(judg)}
if len(J):
    JA = J[J["pass"] == "A"]; fin = JA[JA.doc != "wp29"]
    def paired(a, b, outcome, df, higher_better=True):
        pa = df[df.condition == a].set_index(["task", "doc", "rep"])[outcome]; pb = df[df.condition == b].set_index(["task", "doc", "rep"])[outcome]
        d = (pa - pb).dropna()
        if len(d) < 3: return {"n": int(len(d))}
        cl = d.reset_index().groupby(["task", "doc"])[outcome].mean(); rng = np.random.default_rng(45)
        boots = [float(rng.choice(cl.values, len(cl)).mean()) for _ in range(2000)]
        try: p = float(stats.wilcoxon(d.values, zero_method="wilcox").pvalue) if (d != 0).any() else 1.0
        except Exception: p = None
        x = pa.loc[d.index].values; y = pb.loc[d.index].values; cliff = float(np.mean([np.sign(xi - yi) for xi in x for yi in y]))
        return {"n": int(len(d)), "mean_diff": round(float(d.mean()), 3), "ci95": [round(float(np.percentile(boots, 2.5)), 3), round(float(np.percentile(boots, 97.5)), 3)], "wilcoxon_p": p, "cliffs_delta": round(cliff, 3)}
    H = {"H1_D_vs_B": ("D_profiles_retrieval", "B_strong_prompt", ["clarity", "argument_logic", "contribution_positioning", "usefulness"]), "H2_D_vs_C": ("D_profiles_retrieval", "C_conventional_rag", ["clarity", "argument_logic", "contribution_positioning", "usefulness"]),
         "H3_E_vs_D": ("E_full_engine", "D_profiles_retrieval", ["methodological_rigor", "fidelity", "usefulness"]), "E_vs_A": ("E_full_engine", "A_baseline", DIMS), "E_vs_B": ("E_full_engine", "B_strong_prompt", DIMS)}
    res["hypotheses"] = {}
    for h, (a, b, outs) in H.items():
        tests = {o: paired(a, b, o, fin) for o in outs}; ps = [(o, t["wilcoxon_p"]) for o, t in tests.items() if t.get("wilcoxon_p") is not None]
        m = len(ps); holm = {}
        for rank, (o, p) in enumerate(sorted(ps, key=lambda x: x[1])): holm[o] = min(1.0, p * (m - rank))
        for o in tests:
            if o in holm: tests[o]["holm_p"] = round(holm[o], 4); tests[o]["significant"] = bool(holm[o] < 0.05 and (tests[o]["ci95"][0] > 0 or tests[o]["ci95"][1] < 0))
        res["hypotheses"][h] = tests
    res["ablations"] = {ab: {o: paired("E_full_engine", ab, o, fin) for o in DIMS} for ab in ["E_no_retrieval", "E_no_profile", "E_no_referee", "E_no_factcheck"] if (fin.condition == ab).any()}
    means = fin.groupby("condition")[DIMS + ["n_invented_facts"]].mean().round(2); means["n"] = fin.groupby("condition").size(); means.to_csv(AN / "judge_means_by_condition.csv")
    res["judge_means"] = json.loads(means.to_json(orient="index"))
    res["task_means"] = json.loads(fin.groupby(["task", "condition"])[DIMS].mean().round(2).to_json(orient="index"))
    ood = JA[JA.doc == "wp29"]
    if len(ood): res["out_of_domain_means"] = json.loads(ood.groupby("condition")[DIMS].mean().round(2).to_json(orient="index"))
    # judge consistency
    JB = J[J["pass"] == "B"]
    if len(JB):
        mm = JA.merge(JB, on=["task", "doc", "rep", "condition"], suffixes=("_a", "_b"))
        res["judge_consistency"] = {"n_pairs": int(len(mm)), **{d: {"mean_abs_diff": round(float((mm[d + "_a"] - mm[d + "_b"]).abs().mean()), 3), "spearman": round(float(stats.spearmanr(mm[d + "_a"], mm[d + "_b"]).correlation), 3) if mm[d + "_a"].nunique() > 1 and mm[d + "_b"].nunique() > 1 else None} for d in DIMS}}
    # H4/H5 secondary outcomes
    ok = Rn[Rn.ok & (Rn.domain == "finance")]
    sec = ok.groupby("condition")[["preserved_rate", "new_numbers", "n_unverifiable_citations", "n_unsupported_claims", "max_overlap_share", "todo_placeholders", "n_words", "cost_usd", "latency_s"]].mean().round(3); sec["n"] = ok.groupby("condition").size(); sec.to_csv(AN / "secondary_by_condition.csv")
    res["secondary_means"] = json.loads(sec.to_json(orient="index"))
    res["H4"] = {f"E_vs_{c[0]}": {o: paired("E_full_engine", c, o, ok.assign(rep=ok.rep)) for o in ["n_unverifiable_citations", "n_unsupported_claims", "new_numbers"]} for c in ["A_baseline", "B_strong_prompt", "C_conventional_rag"]}
    res["H5_overlap"] = {"max_overlap_share_overall": float(ok.max_overlap_share.max()), "share_outputs_below_0.02": float((ok.max_overlap_share < 0.02).mean())}
    # LaTeX tables
    def tex_table(df, caption, label, fname):
        cols = list(df.columns); L = ["\\begin{table}[htbp]\\centering\\footnotesize", f"\\caption{{{caption}}}\\label{{{label}}}", "\\begin{tabular}{l" + "r" * len(cols) + "}\\toprule", "Condition & " + " & ".join(c.replace("_", "\\_") for c in cols) + " \\\\\\midrule"]
        L += [f"{idx.replace('_', '\\_')} & " + " & ".join(("" if pd.isna(v) else (f"{v:.2f}" if isinstance(v, float) else str(v))) for v in row) + " \\\\" for idx, row in df.iterrows()]
        L += ["\\bottomrule\\end{tabular}", "\\begin{minipage}{0.95\\linewidth}\\vspace{2pt}\\scriptsize Notes: generated by scripts/05\\_analyze\\_benchmark.py from recorded runs (protocol heldout\\_v1); judge = claude-sonnet-5-5 blind to condition; scores 1--7.\\end{minipage}", "\\end{table}"]
        (TAB / fname).write_text("\n".join(L) + "\n")
    tex_table(means, "Judge scores by condition (finance held-out documents, pass A)", "tab:judge_means", "tab_judge_means.tex")
    tex_table(sec, "Deterministic secondary outcomes by condition (finance held-out documents)", "tab:secondary", "tab_secondary.tex")
    hyp_rows = []
    for h, tests in res["hypotheses"].items():
        for o, t in tests.items():
            if "mean_diff" in t: hyp_rows.append({"contrast": h, "outcome": o, "n": t["n"], "mean_diff": t["mean_diff"], "ci_low": t["ci95"][0], "ci_high": t["ci95"][1], "holm_p": t.get("holm_p"), "cliffs_delta": t["cliffs_delta"]})
    hd = pd.DataFrame(hyp_rows); hd.to_csv(AN / "hypothesis_tests.csv", index=False)
    if len(hd): tex_table(hd.set_index(hd.contrast + " / " + hd.outcome).drop(columns=["contrast", "outcome"]), "Preregistered paired contrasts (cluster bootstrap 95\\% CI, Holm-adjusted Wilcoxon p)", "tab:hypotheses", "tab_hypotheses.tex")
    # figure
    viz.style(); import matplotlib.pyplot as plt
    conds = [c for c in ["A_baseline", "B_strong_prompt", "C_conventional_rag", "D_profiles_retrieval", "E_full_engine"] if c in means.index]
    fig, ax = plt.subplots(figsize=(7, 3)); w = 0.8 / len(DIMS)
    for i, d in enumerate(DIMS):
        vals = [means.loc[c, d] for c in conds]; se = [fin[fin.condition == c][d].std() / np.sqrt(max((fin.condition == c).sum(), 1)) for c in conds]
        ax.bar(np.arange(len(conds)) + i * w, vals, w, yerr=[1.96 * s for s in se], label=d, color=viz.PALETTE[i], edgecolor=viz.SURFACE, linewidth=0.8, error_kw={"elinewidth": 0.8, "ecolor": viz.TEXT2})
    ax.set_xticks(np.arange(len(conds)) + w * (len(DIMS) - 1) / 2); ax.set_xticklabels([c.split("_")[0] for c in conds]); ax.set_ylim(1, 7); ax.set_ylabel("Judge score (1–7)"); ax.legend(fontsize=6, ncol=3); ax.set_title("Blind judge scores by condition (finance held-out tasks; ±1.96 SE)", fontsize=9)
    fig.tight_layout(); fig.savefig(R / "paper" / "figures" / "fig_judge_by_condition.pdf"); fig.savefig(R / "figures" / "fig_judge_by_condition.pdf"); plt.close(fig)
(AN / f"results_{PROTO}.json").write_text(json.dumps(res, indent=1, default=str)); print(json.dumps({k: v for k, v in res.items() if k in ("n_runs", "n_runs_ok", "n_judgments", "hypotheses", "judge_means", "H5_overlap", "judge_consistency")}, indent=1, default=str)[:6000])