#!/usr/bin/env python3
"""exp15 — Q0 RE-ANALYSIS from existing judge_results_Q0.json (no re-run; same flat-chunking-both-sides
reasoning as Q3a). Applies the orchestrator's corrections:

 - TOST: the +0.041 point estimate is correct, but equivalence is NOT established. 90% CI upper bound
   exceeds the margin M=0.05 -> by the freeze rule the verdict is "insufficient data", and the §D5
   branch "editor mandates revised" does NOT trigger. (a0 = do-nothing DRAFT, arm = A0 editor FINAL;
   paired_diff = editor - donothing, positive = editor helps.)
 - SYMMETRY: the editor CUTS catastrophe votes (donothing -> editor), an effect ~the same size as the
   Q3a carryover "plus" — must sit in the headline next to it.
 - UNITS: catastrophe counts are judge VOTES, not traps. Report the distinct-trap counts and which traps
   the editor added flags on vs cleared.

Run: python q0_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")
M = 0.05   # TOST equivalence margin (freeze §1.5)


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


def tost(xs, margin):
    """Two one-sided tests for equivalence within +-margin. Equivalent iff BOTH p<0.05 (== 90% CI inside
    (-margin, +margin))."""
    xs = np.asarray(xs, float); n = len(xs)
    m = xs.mean(); se = xs.std(ddof=1) / math.sqrt(n)
    if se == 0:
        return None
    t_lo = (m - (-margin)) / se       # H0: mu <= -margin
    t_hi = (m - margin) / se          # H0: mu >= +margin
    p_lo = 1 - stats.t.cdf(t_lo, n - 1)
    p_hi = stats.t.cdf(t_hi, n - 1)
    return {"p_lower": round(float(p_lo), 4), "p_upper": round(float(p_hi), 4),
            "equivalent_90": bool(p_lo < 0.05 and p_hi < 0.05)}


def main():
    d = json.load(open(SD / "judge_results_Q0.json"))
    res = d["results"] if isinstance(d, dict) and "results" in d else d
    allr = res
    prim = [r for r in res if r["type"] in PRIMARY]

    out = {}
    for label, rows in (("all", allr), ("primary", prim), ("T-ana", [r for r in res if r["type"] == "T-ana"])):
        pd = [r["paired_diff_arm_minus_a0"] for r in rows if r["paired_diff_arm_minus_a0"] is not None]
        m, sd, lo, hi, n = ci90(pd)
        t = tost(pd, M)
        # catastrophe: arm = editor(final), a0 = donothing(draft)
        cat_editor = sum(r["catastrophe"]["arm"] for r in rows)
        cat_donothing = sum(r["catastrophe"]["a0"] for r in rows)
        out[label] = {"editor_minus_donothing_mean": m, "sd": sd, "CI90": [lo, hi], "n": n,
                      "TOST_M": M, "TOST": t,
                      "equivalence_established": bool(t and t["equivalent_90"]),
                      "cat_votes_editor_vs_donothing": [cat_editor, cat_donothing],
                      "cat_reduction_pct": round(100 * (cat_donothing - cat_editor) / cat_donothing, 1) if cat_donothing else None}
        verdict = ("EQUIVALENT (CI within +-M)" if (t and t["equivalent_90"])
                   else f"NOT established (CI upper {hi} exceeds M={M}) -> insufficient data, D5 branch NOT triggered")
        print(f"[{label}] editor-donothing {m}  sd {sd}  CI90 [{lo},{hi}] n={n}")
        print(f"        TOST {t}  -> {verdict}")
        print(f"        catastrophe VOTES editor {cat_editor} vs donothing {cat_donothing} "
              f"(editor cuts {out[label]['cat_reduction_pct']}%)")

    # UNITS: distinct traps with catastrophe votes (not the same as vote counts); who added/cleared flags
    added, cleared, net = [], [], []
    for r in prim:
        ce, cd = r["catastrophe"]["arm"], r["catastrophe"]["a0"]   # editor, donothing
        if ce > cd:
            added.append((r["trap"], cd, ce))
        elif ce < cd:
            cleared.append((r["trap"], cd, ce))
    trap_editor = sorted({r["trap"] for r in prim if r["catastrophe"]["arm"] > 0})
    trap_donothing = sorted({r["trap"] for r in prim if r["catastrophe"]["a0"] > 0})
    out["units"] = {"distinct_traps_with_editor_cat": len(trap_editor),
                    "distinct_traps_with_donothing_cat": len(trap_donothing),
                    "editor_ADDED_flags_on": added, "editor_CLEARED_flags_on": cleared}
    print(f"\n[units] distinct traps w/ catastrophe: editor {len(trap_editor)} vs donothing {len(trap_donothing)} "
          f"(votes != traps)")
    print(f"  editor ADDED flags on {len(added)} traps: {added}")
    print(f"  editor CLEARED flags on {len(cleared)} traps: {cleared}")

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


if __name__ == "__main__":
    main()
