#!/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"]}
    # H5: the preregistered metric (max 8-gram overlap with ANY corpus document) counts overlap with the task's own source document
    # (e.g. improve_abstract necessarily reuses the source abstract). Post-hoc refinement, reported as such: recompute excluding the source document.
    from jof250 import auditing
    other = []
    for r in runs:
        if not r.get("ok") or r["domain"] != "finance": continue
        ov = auditing.overlap_with_corpus(r["text"]); oth = [d for d in ov["top_docs"] if d["doc_id"] != r["doc"]]
        other.append({"task": r["task"], "doc": r["doc"], "condition": r["condition"], "rep": r["rep"], "self_share": next((d["share"] for d in ov["top_docs"] if d["doc_id"] == r["doc"]), 0.0), "other_share": oth[0]["share"] if oth else 0.0})
    OV = pd.DataFrame(other); OV.to_csv(AN / "overlap_by_run.csv", index=False)
    res["H5_overlap"] = {"preregistered_metric": {"max_overlap_share_overall": float(ok.max_overlap_share.max()), "share_outputs_below_0.02": float((ok.max_overlap_share < 0.02).mean())},
                         "excluding_source_document": {"max_other_share": float(OV.other_share.max()), "share_outputs_below_0.02": float((OV.other_share < 0.02).mean()), "mean_other_share": float(OV.other_share.mean())},
                         "self_overlap_by_task": json.loads(OV.groupby("task").self_share.mean().round(3).to_json()), "self_overlap_by_condition": json.loads(OV.groupby("condition").self_share.mean().round(3).to_json())}
    # 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])