#!/usr/bin/env python3
"""exp15 — Q3a RE-ANALYSIS from existing judge_results_Q3a.json (no re-run: the head-excerpt artifact is
Q1b-specific — Q3a's a0 and every arm share FLAT chunking, so the dependent sits at the head in BOTH
sides -> no truncation asymmetry). Applies the orchestrator's corrections:

 (b) primary = T-ana + T-ell (pre-reg §1.5), NOT the post-hoc T-ana-only "+0.042"; disclose T-ell-43.
 (c) BOTH estimators reported with 90% CI: paired (mean per-cell arm-a0) AND marginal arm-acc
     (mean arm_acc - mean a0_acc). They diverge; report both.
 (d) leave-one-judge-out sign per mode (src stable? target-side/final flip?).
 (e) POOLED fix-rate: A0-failed traps defined ONCE by pooling a0 across modes (the per-mode denominators
     1/4,1/4,0/3 were an artifact of re-judging A0 separately per mode). fix = arm resolves it.
 (a) a0<->a0 NOISE PROXY: the same flat-a0 chunk is judged 3x (once per mode-cell) -> its a0_accuracy
     spread across modes is pure judge noise. Reported as a preliminary floor; the dedicated FLOOR block
     (A0 vs A0', judge_run.py --run FLOOR) is authoritative.

Run: python q3a_reanalysis.py
"""
from __future__ import annotations
import json
import math
from collections import defaultdict
from pathlib import Path

import numpy as np
from scipy import stats

SD = Path("/home/ubuntu/books/gu-zhenren/exp15")
PRIMARY = ("T-ana", "T-ell")
MODES = ("q3a-src", "q3a-draft", "q3a-final")
FAIL = 0.5   # a0_accuracy < FAIL  == "A0 failed this trap"  (correct=1/partial=.5/wrong=0)


def ci90(xs):
    xs = np.asarray(xs, float)
    n = len(xs)
    if n < 2:
        return (float(xs.mean()) if n else 0.0, float("nan"), float("nan"), n)
    m = xs.mean(); se = xs.std(ddof=1) / math.sqrt(n)
    t = stats.t.ppf(0.95, n - 1)
    return round(m, 4), round(m - t * se, 4), round(m + t * se, 4), n


def paired_p(xs):
    xs = np.asarray(xs, float)
    if len(xs) < 2 or xs.std() == 0:
        return None
    return round(float(stats.ttest_1samp(xs, 0.0).pvalue), 3)


def main():
    d = json.load(open(SD / "judge_results_Q3a.json"))
    res = d["results"] if isinstance(d, dict) and "results" in d else d
    by_mode = defaultdict(list)
    for r in res:
        by_mode[r["arm"]].append(r)

    # a0<->a0 noise proxy: spread of a0_accuracy for the same trap across its mode-cells (identical a0 text)
    a0_by_trap = defaultdict(dict)
    for r in res:
        if r["type"] in PRIMARY:
            a0_by_trap[r["trap"]][r["arm"]] = r["a0_accuracy"]
    noises = []
    for trap, mm in a0_by_trap.items():
        vals = [v for v in mm.values() if v is not None]
        if len(vals) >= 2:
            noises.append(max(vals) - min(vals))
    floor_proxy = {"mean_a0_spread": round(float(np.mean(noises)), 4),
                   "max_a0_spread": round(float(np.max(noises)), 4), "n_traps": len(noises)}

    out = {"floor_proxy_a0_vs_a0": floor_proxy, "modes": {}}
    print(f"[a0<->a0 noise proxy] mean spread {floor_proxy['mean_a0_spread']}  "
          f"max {floor_proxy['max_a0_spread']}  (dedicated FLOOR block is authoritative)")

    for mode in MODES:
        rows = by_mode.get(mode, [])
        prim = [r for r in rows if r["type"] in PRIMARY]
        tana = [r for r in rows if r["type"] == "T-ana"]
        # paired estimator (per-cell arm-a0)
        pd_prim = [r["paired_diff_arm_minus_a0"] for r in prim if r["paired_diff_arm_minus_a0"] is not None]
        pd_tana = [r["paired_diff_arm_minus_a0"] for r in tana if r["paired_diff_arm_minus_a0"] is not None]
        # marginal arm-acc estimator (mean arm - mean a0)
        arm_prim = np.mean([r["arm_accuracy"] for r in prim])
        a0_prim = np.mean([r["a0_accuracy"] for r in prim])
        marg_prim = round(float(arm_prim - a0_prim), 4)
        # LOO sign (mean of per-cell LOO diffs, which drop one judge)
        loo = {"without_grok-4.3": [], "without_mistral-large-latest": []}
        for r in prim:
            for k, v in (r.get("leave_one_judge_out") or {}).items():
                if k in loo:
                    loo[k].append(v)
        loo_means = {k: round(float(np.mean(v)), 4) for k, v in loo.items() if v}
        # catastrophes
        cat_arm = sum(r["catastrophe"]["arm"] for r in prim)
        cat_a0 = sum(r["catastrophe"]["a0"] for r in prim)
        pm, plo, phi, pn = ci90(pd_prim)
        tm, tlo, thi, tn = ci90(pd_tana)
        mde = round(2.8 * np.std(pd_prim, ddof=1) / math.sqrt(len(pd_prim)), 4) if len(pd_prim) > 1 else None
        out["modes"][mode] = {
            "primary_paired_mean": pm, "primary_paired_CI90": [plo, phi], "primary_n": pn,
            "primary_paired_p": paired_p(pd_prim), "primary_MDE90": mde,
            "primary_marginal_armacc": marg_prim,
            "tana_only_paired_mean": tm, "tana_only_CI90": [tlo, thi], "tana_n": tn,
            "loo_mean_diff": loo_means,
            "primary_catastrophes_arm_vs_a0": [cat_arm, cat_a0],
        }
        print(f"\n[{mode}]")
        print(f"  primary(T-ana+T-ell) PAIRED {pm}  CI90 [{plo},{phi}] n={pn} p={paired_p(pd_prim)} MDE90~{mde}")
        print(f"  primary MARGINAL arm-acc      {marg_prim}   (metric-swap gap vs paired = {round(marg_prim-pm,4)})")
        print(f"  T-ana-only PAIRED             {tm}  CI90 [{tlo},{thi}] n={tn}   <- the reported headline")
        print(f"  LOO mean diff {loo_means}   catastrophes arm/a0 {cat_arm}/{cat_a0}")

    # pooled fix-rate: A0-failed defined by pooling a0 across modes (mean a0 over the 3 cells < FAIL)
    pooled_a0 = {t: np.mean(list(v.values())) for t, v in a0_by_trap.items()}
    a0_failed = [t for t, a in pooled_a0.items() if a < FAIL]
    fixrate = {}
    for mode in MODES:
        rows = {r["trap"]: r for r in by_mode.get(mode, []) if r["type"] in PRIMARY}
        fixed = sum(1 for t in a0_failed if t in rows and rows[t]["arm_accuracy"] >= FAIL)
        fixrate[mode] = f"{fixed}/{len(a0_failed)} ({round(100*fixed/len(a0_failed)) if a0_failed else 0}%)"
    out["pooled_fix_rate"] = {"A0_failed_traps": sorted(a0_failed), "n_failed": len(a0_failed), "by_mode": fixrate}
    print(f"\n[pooled fix-rate] A0-failed (pooled a0<{FAIL}) = {len(a0_failed)} traps {sorted(a0_failed)}")
    for mode, fr in fixrate.items():
        print(f"  {mode}: {fr}")

    # T-ell-43 disclosure (winner-caused catastrophe)
    tell43 = {r["arm"]: {"a0_acc": r["a0_accuracy"], "arm_acc": r["arm_accuracy"],
                         "cat_arm": r["catastrophe"]["arm"], "cat_a0": r["catastrophe"]["a0"]}
              for r in res if r["trap"] == "T-ell-43"}
    out["T-ell-43_disclosure"] = tell43
    print(f"\n[T-ell-43 disclosure] {json.dumps(tell43, ensure_ascii=False)}")

    (SD / "q3a_reanalysis.json").write_text(json.dumps(out, ensure_ascii=False, indent=2), encoding="utf-8")
    print(f"\n-> {SD/'q3a_reanalysis.json'}")


if __name__ == "__main__":
    main()
