#!/usr/bin/env python3
"""exp15 — Q1b (full-chunk re-judge) + FLOOR analysis. Reads the fresh full-chunk judge_results_Q1b.json
and judge_results_FLOOR.json and produces the corrected, floor-relative Q1b verdict.

Key checks (the whole point of the re-judge):
 1. CORRECTED effect: oracle-minus-flat paired diff, primary (T-ana+T-ell) and per type, with 90% CI, LOO,
    catastrophes — vs the v1 head-excerpt headline (-0.564) to show the artifact's magnitude.
 2. ARTIFACT-GONE check: split by oracle_frac (dep position in the oracle chunk). Under the old 1100-char
    head-excerpt, dep-in-window cells were +0.238 and dep-out were -0.872 (the truncation signature). With
    FULL chunks the split MUST collapse (both similar) — if dep-out is still strongly negative the fix failed.
 3. FLOOR-relative: is |Q1b effect| within the A0-vs-A0' full-chunk noise floor? Effects under the floor
    are not interpretable (freeze §1.5).
 4. Per-vote audit available (judges.py now persists votes) — catastrophe counts reconstructable.

Run (after both blocks): python q1b_reanalysis.py
"""
from __future__ import annotations
import json
import math
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")


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


def load(name):
    p = SD / name
    if not p.exists():
        return None
    d = json.load(open(p))
    return d["results"] if isinstance(d, dict) and "results" in d else d


def block_stats(rows, label):
    pd = [r["paired_diff_arm_minus_a0"] for r in rows if r.get("paired_diff_arm_minus_a0") is not None]
    m, lo, hi, n = ci90(pd)
    cat_arm = sum(r["catastrophe"]["arm"] for r in rows)
    cat_a0 = sum(r["catastrophe"]["a0"] for r in rows)
    loo = {"without_grok-4.3": [], "without_mistral-large-latest": []}
    for r in rows:
        for k, v in (r.get("leave_one_judge_out") or {}).items():
            if k in loo:
                loo[k].append(v)
    loo_m = {k: round(float(np.mean(v)), 4) for k, v in loo.items() if v}
    trunc = sum(1 for r in rows if r.get("budget_truncated"))
    return {"label": label, "mean": m, "CI90": [lo, hi], "n": n,
            "catastrophes_arm_vs_a0": [cat_arm, cat_a0], "loo_mean": loo_m, "budget_truncated_cells": trunc}


def main():
    q1b = load("judge_results_Q1b.json")
    floor = load("judge_results_FLOOR.json")
    v1 = load("judge_results_Q1b_v1_headexcerpt.json")
    out = {}

    if q1b is None:
        print("Q1b results not present yet."); return
    prim = [r for r in q1b if r["type"] in PRIMARY]
    tana = [r for r in q1b if r["type"] == "T-ana"]
    tent = [r for r in q1b if r["type"] == "T-ent"]
    out["Q1b_fullchunk"] = {"primary": block_stats(prim, "primary(T-ana+T-ell)"),
                            "T-ana": block_stats(tana, "T-ana"),
                            "T-ent_control": block_stats(tent, "T-ent control") if tent else None}
    print("=== Q1b FULL-CHUNK (oracle - flat; + = oracle better) ===")
    for k in ("primary", "T-ana", "T-ent_control"):
        s = out["Q1b_fullchunk"][k]
        if s:
            print(f"  {s['label']:22} mean {s['mean']:+.4f}  CI90 {s['CI90']}  n={s['n']}  "
                  f"cat oracle/flat {s['catastrophes_arm_vs_a0']}  LOO {s['loo_mean']}  trunc={s['budget_truncated_cells']}")

    # artifact-gone check: split by oracle_frac
    din = [r for r in prim if (r.get("oracle_frac") or 0) < 0.5]
    dout = [r for r in prim if (r.get("oracle_frac") or 0) >= 0.5]
    si, so = ci90([r["paired_diff_arm_minus_a0"] for r in din]), ci90([r["paired_diff_arm_minus_a0"] for r in dout])
    out["artifact_gone_check"] = {"dep_in_window_frac<0.5": {"mean": si[0], "n": si[3]},
                                  "dep_out_window_frac>=0.5": {"mean": so[0], "n": so[3]},
                                  "v1_reference": {"dep_in": 0.238, "dep_out": -0.872},
                                  "collapsed": abs(si[0] - so[0]) < 0.3}
    print(f"\n=== ARTIFACT-GONE check (full chunks should collapse the split) ===")
    print(f"  dep-in  (ofrac<0.5): mean {si[0]:+.4f} n={si[3]}   (v1 head-excerpt: +0.238)")
    print(f"  dep-out (ofrac>=0.5): mean {so[0]:+.4f} n={so[3]}   (v1 head-excerpt: -0.872)")
    print(f"  -> split {'COLLAPSED (artifact gone)' if out['artifact_gone_check']['collapsed'] else 'STILL PRESENT — investigate'}")

    if v1:
        vp = [r for r in v1 if r["type"] == "T-ana"]
        vm, *_ = ci90([r["paired_diff_arm_minus_a0"] for r in vp])
        out["v1_vs_v2"] = {"v1_Tana_headexcerpt": vm, "v2_Tana_fullchunk": out["Q1b_fullchunk"]["T-ana"]["mean"],
                           "artifact_magnitude": round(out["Q1b_fullchunk"]["T-ana"]["mean"] - vm, 4)}
        print(f"\n=== v1 (head-excerpt) vs v2 (full-chunk) T-ana ===")
        print(f"  v1 {vm:+.4f}  ->  v2 {out['Q1b_fullchunk']['T-ana']['mean']:+.4f}  "
              f"(the artifact moved the headline by {out['v1_vs_v2']['artifact_magnitude']:+.4f})")

    if floor is not None:
        fpd = [abs(r["paired_diff_arm_minus_a0"]) for r in floor if r.get("paired_diff_arm_minus_a0") is not None]
        fmean = round(float(np.mean(fpd)), 4); fmax = round(float(np.max(fpd)), 4)
        fsigned = ci90([r["paired_diff_arm_minus_a0"] for r in floor if r.get("paired_diff_arm_minus_a0") is not None])
        out["FLOOR_fullchunk"] = {"mean_abs_diff": fmean, "max_abs_diff": fmax,
                                  "signed_mean": fsigned[0], "signed_CI90": [fsigned[1], fsigned[2]], "n": fsigned[3]}
        prim_abs = abs(out["Q1b_fullchunk"]["primary"]["mean"])
        out["FLOOR_fullchunk"]["Q1b_primary_exceeds_floor"] = bool(prim_abs > fmean)
        print(f"\n=== FLOOR (A0 vs A0', full-chunk) ===")
        print(f"  |diff| mean {fmean}  max {fmax}  signed {fsigned[0]:+.4f} CI90 [{fsigned[1]},{fsigned[2]}] n={fsigned[3]}")
        print(f"  Q1b primary |{prim_abs:.4f}| {'EXCEEDS' if prim_abs>fmean else 'is WITHIN'} the floor {fmean}")
    else:
        print("\n(FLOOR results not present yet — run FLOOR block first.)")

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


if __name__ == "__main__":
    main()