#!/usr/bin/env python3
"""exp15/Q4a — DETERMINISTIC scorer (PRIMARY instrument, $0). Reads the paired 2x2 cells generated
by q4a_generate.py, DET-scores each meaning-trap with the exp14b rules (q4a_traps), and produces:

  * 2x2 fix-rate matrix   {flash,pro} x {raw, +glm-editor}
  * §D5 primary rule      among traps FLASH-raw FAILS, fraction PRO-raw FIXES (threshold >=0.5)
  * editor delta          (pro_glm - pro_raw) and (flash_glm - flash_raw): does the editor
                          PRESERVE or BREAK the draft's fidelity (T-inv = editor property, D38)
  * catastrophe class     a1/a2/c1 (polarity/rank inversions) tracked separately
  * seed variance         flash-s0 vs flash-s1 (and pro) disagreement = the DET-level noise floor
                          (should be ~0 for polarity/number traps -> confirms instrument reliability)
  * provenance x-check    DET-score the on-disk BACKEND flash draft (records.json) + exp14b arm C,
                          to confirm the freshly-regenerated flash side agrees with the saved one
  * COGS delta            pro-draft vs flash-draft stage $ from the gen ledger (metric iii; pro ~x3.1)

Scoring convention: fix=1, fail=0, '?'=undecidable (EXCLUDED from rates, reported separately —
addendum §2: no silent drops). Cell score over seeds = mean of decided seeds; cell 'fails' if
score<0.5, 'fixes' if score>=0.5, 'undecided' if all seeds '?'. A cell whose file(s) are MISSING
(not generated / valley-defer / GEN hard-stop) is tracked as 'missing' SEPARATELY from all-'?'
(review F1). §D5 ratify/not-ratify is gated on (i) a complete 2x2 square for the trap and (ii) a
pre-registered minimum decided denominator (>=3); otherwise the verdict is INSUFFICIENT-POWER, never
a clean pass/fail off a tiny partial set. $0 (no api calls)."""
from __future__ import annotations
import json
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
import q4a_traps as Q

SD = Path("/home/ubuntu/books/gu-zhenren/exp15")
OUT = SD / "q4a"
GEN_LEDGER = SD / "q4a_gen_costs.jsonl"
CELLS = ["flash_raw", "pro_raw", "flash_glm", "pro_glm"]
SEEDS = [0, 1, 2]
RECS = {(r["chapter"], r["chunk_idx"]): r
        for r in json.load(open("/home/ubuntu/books/gu-zhenren/rerun/records.json"))}
EX14B = Path("/home/ubuntu/books/gu-zhenren/exp14b/arms")


def _read(cell, cid, seed):
    p = OUT / cell / f"{cid}.s{seed}.txt"
    return p.read_text(encoding="utf-8") if p.exists() else None


def _cell_verdict(scores):
    """scores = list of 'fix'/'fail'/'?' over seeds -> (label, mean_of_decided|None, n_decided, n_q)."""
    dec = [1.0 if s == "fix" else 0.0 for s in scores if s in ("fix", "fail")]
    nq = sum(1 for s in scores if s == "?")
    if not dec:
        return "undecided", None, 0, nq
    m = sum(dec) / len(dec)
    return ("fix" if m >= 0.5 else "fail"), round(m, 3), len(dec), nq


def score_all():
    traps = Q.traps()
    per_trap = []
    for t in traps:
        row = {"id": t["id"], "chunk": t["chunk"], "cat_class": t["catastrophe_class"], "desc": t["desc"],
               "cells": {}, "seed_raw": {}}
        row["seed_located"] = {}
        for cell in CELLS:
            seed_scores, located = [], []
            for seed in SEEDS:
                txt = _read(cell, t["chunk"], seed)
                if txt is None:
                    seed_scores.append(None); located.append(None); continue
                v, span = Q.det_score(txt, t)
                seed_scores.append(v)
                located.append(bool(span))     # sentence found = phenomenon RENDERED (else omitted/locator-miss)
            row["seed_raw"][cell] = seed_scores
            row["seed_located"][cell] = located
            present = [s for s in seed_scores if s is not None]
            label, mean, ndec, nq = _cell_verdict(present)
            missing = len(present) < len(SEEDS)          # some seed file absent (not generated / hard-stop)
            if len(present) == 0:
                label = "missing"
            loc = [x for x in located if x is not None]
            row["cells"][cell] = {"label": label, "mean": mean, "n_decided": ndec, "n_q": nq,
                                  "n_present": len(present), "missing": missing,
                                  "n_located": sum(1 for x in loc if x), "n_seen": len(loc)}
        # a trap is COMPLETE only if every cell has all seeds present (review F1: completeness gate)
        row["complete"] = all(not row["cells"][c]["missing"] for c in CELLS)
        row["paired_decided"] = all(row["cells"][c]["mean"] is not None for c in CELLS)
        per_trap.append(row)
    return traps, per_trap


MIN_DENOM = 3   # pre-registered minimum decided flash-fail denominator for a §D5 ratify/not-ratify verdict


def cell_fixrate(per_trap, cell, cat_only=False, paired_only=False):
    """MARGINAL fix-rate: over this cell's own decided traps (paired_only=True -> only traps decided in
    ALL 4 cells, so the four cells are comparable — the paired headline, review F6)."""
    rows = [r for r in per_trap if (r["cat_class"] if cat_only else True)
            and (r["paired_decided"] if paired_only else True)]
    dec = [r["cells"][cell]["mean"] for r in rows if r["cells"][cell]["mean"] is not None]
    return (round(sum(1 for m in dec if m >= 0.5) / len(dec), 3), len(dec)) if dec else (None, 0)


def d5_rule(per_trap):
    """§D5 (review F1): among COMPLETE traps FLASH-raw FAILS (mean<0.5), fraction PRO-raw FIXES
    (mean>=0.5). ratified is a clean bool ONLY when the denominator >= MIN_DENOM AND the 2x2 square is
    complete for every counted trap; otherwise ratified=None and status='INSUFFICIENT-POWER'."""
    flash_fail, pro_fixes, detail = [], 0, []
    incomplete = [r["id"] for r in per_trap if not r["complete"]]
    for r in per_trap:
        if not r["complete"]:
            continue                      # never let a partial square enter the ratification metric
        f = r["cells"]["flash_raw"]["mean"]; p = r["cells"]["pro_raw"]["mean"]
        if f is None or p is None:
            continue
        if f < 0.5:                       # flash-raw fails this meaning-trap
            flash_fail.append(r["id"])
            fixed = p >= 0.5
            pro_fixes += int(fixed)
            detail.append({"id": r["id"], "cat": r["cat_class"], "flash": f, "pro": p, "pro_fixes": fixed})
    n = len(flash_fail)
    rate = round(pro_fixes / n, 3) if n else None
    powered = n >= MIN_DENOM
    ratified = (rate >= 0.5) if (powered and rate is not None) else None
    status = ("RATIFIED" if ratified else "NOT-RATIFIED") if powered else "INSUFFICIENT-POWER"
    return {"flash_fail_traps": flash_fail, "n_denominator": n, "min_denominator": MIN_DENOM,
            "pro_fixes": pro_fixes, "fix_rate": rate, "threshold": 0.5, "powered": powered,
            "ratified": ratified, "status": status, "incomplete_traps": incomplete, "detail": detail}


def editor_delta(per_trap, side):
    """(side_glm mean - side_raw mean) per trap; positive = editor improves, negative = editor breaks."""
    deltas, breaks = [], []
    for r in per_trap:
        raw = r["cells"][f"{side}_raw"]["mean"]; glm = r["cells"][f"{side}_glm"]["mean"]
        if raw is None or glm is None:
            continue
        d = glm - raw
        deltas.append(d)
        if d < 0:
            breaks.append({"id": r["id"], "raw": raw, "glm": glm})
    mean = round(sum(deltas) / len(deltas), 3) if deltas else None
    return {"mean_delta": mean, "n": len(deltas), "editor_breaks": breaks}


def seed_floor(per_trap):
    """DET-level noise: per cell, fraction of traps whose decided seeds are NOT unanimous on fix/fail
    (>=2 decided seeds required). ~0 expected for polarity/number traps -> confirms DET reliability."""
    out = {}
    for cell in CELLS:
        disagree, n = 0, 0
        for r in per_trap:
            ss = [s for s in r["seed_raw"][cell] if s in ("fix", "fail")]
            if len(ss) >= 2:
                n += 1; disagree += int(len(set(ss)) > 1)
        out[cell] = {"disagree": disagree, "n_multi_decided": n,
                     "rate": round(disagree / n, 3) if n else None}
    return out


def coverage(per_trap):
    """OMISSION metric (sanity-gate finding): per cell, fraction of traps whose phenomenon sentence is
    LOCATED (rendered) vs not (omitted or locator-miss). flash_raw vs pro_raw = does the strong draft
    omit MORE? Reported honestly with the caveat that 'not-located' conflates true omission with a
    locator gap (symmetric-ish across models)."""
    out = {}
    for cell in CELLS:
        loc = tot = 0
        for r in per_trap:
            for x in r["seed_located"][cell]:
                if x is not None:
                    tot += 1; loc += int(x)
        out[cell] = {"located": loc, "seen": tot, "rate": round(loc / tot, 3) if tot else None}
    return out


def editor_restoration(per_trap):
    """Sanity-gate signal: per side, count (trap,seed) where the RAW draft did NOT render the phenomenon
    (not located) but the glm editor DID (located) — the editor, seeing the source, RESTORED an omission.
    Direct evidence the editor stays load-bearing (argues against strong-draft+do-nothing)."""
    out = {}
    for side in ("flash", "pro"):
        restored, broke, both_absent = [], [], 0
        for r in per_trap:
            rawloc = r["seed_located"][f"{side}_raw"]; glmloc = r["seed_located"][f"{side}_glm"]
            for i in SEEDS:
                rw, gl = rawloc[i], glmloc[i]
                if rw is None or gl is None:
                    continue
                if not rw and gl:
                    restored.append({"id": r["id"], "seed": i})
                elif rw and not gl:
                    broke.append({"id": r["id"], "seed": i})
                elif not rw and not gl:
                    both_absent += 1
        out[side] = {"restored": restored, "n_restored": len(restored),
                     "broke": broke, "n_broke": len(broke), "both_absent": both_absent}
    return out


def provenance_xcheck(traps):
    """DET-score the SAVED backend flash artifacts to confirm the fresh flash side agrees."""
    out = {"backend_flash_raw": {}, "exp14b_arm_C": {}}
    for t in traps:
        ch, ck = t["chunk"].split("."); r = RECS.get((int(ch), int(ck)))
        v_bk, _ = Q.det_score(r["draft"] if r else "", t)
        out["backend_flash_raw"][t["id"]] = v_bk
        cp = EX14B / "C" / f"{t['chunk']}.txt"
        v_c, _ = Q.det_score(cp.read_text(encoding="utf-8") if cp.exists() else "", t)
        out["exp14b_arm_C"][t["id"]] = v_c
    return out


def cogs():
    if not GEN_LEDGER.exists():
        return {}
    agg = {}
    for line in GEN_LEDGER.read_text(encoding="utf-8").splitlines():
        try:
            d = json.loads(line)
        except Exception:
            continue
        cell = d.get("cell", "?"); agg.setdefault(cell, {"cost": 0.0, "n": 0})
        agg[cell]["cost"] += d.get("cost", 0.0); agg[cell]["n"] += 1
    fr = agg.get("flash_raw", {}).get("cost", 0.0); pr = agg.get("pro_raw", {}).get("cost", 0.0)
    return {"by_cell": {k: {"cost": round(v["cost"], 5), "n": v["n"]} for k, v in agg.items()},
            "draft_stage_flash": round(fr, 5), "draft_stage_pro": round(pr, 5),
            "pro_over_flash_x": round(pr / fr, 2) if fr else None}


def main():
    traps, per_trap = score_all()
    # MARGINAL matrix (each cell over its own decided traps) + PAIRED matrix (only traps decided in ALL
    # 4 cells -> the comparable headline; review F6). Marginal cells can be over DIFFERENT trap subsets.
    matrix = {cell: {"marginal": cell_fixrate(per_trap, cell),
                     "paired": cell_fixrate(per_trap, cell, paired_only=True),
                     "cat": cell_fixrate(per_trap, cell, cat_only=True)} for cell in CELLS}
    n_paired = sum(1 for r in per_trap if r["paired_decided"])
    n_complete = sum(1 for r in per_trap if r["complete"])
    result = {
        "n_traps": len(traps), "seeds": SEEDS,
        "n_complete_square": n_complete, "n_paired_decided": n_paired,
        "fixrate_matrix": matrix,
        "d5_rule": d5_rule(per_trap),
        "editor_delta_pro": editor_delta(per_trap, "pro"),
        "editor_delta_flash": editor_delta(per_trap, "flash"),
        "seed_floor": seed_floor(per_trap),
        "coverage": coverage(per_trap),
        "editor_restoration": editor_restoration(per_trap),
        "provenance_xcheck": provenance_xcheck(traps),
        "cogs": cogs(),
        "per_trap": per_trap,
    }
    (SD / "q4a_det_results.json").write_text(json.dumps(result, ensure_ascii=False, indent=2), encoding="utf-8")

    print(f"=== Q4a DET 2x2 fix-rate matrix (n_traps={result['n_traps']}, complete-square={n_complete}, "
          f"paired-decided={n_paired}) ===")
    print("cell = PAIRED (all-4-cells-decided; comparable) | marginal (own decided subset):")
    print(f"{'':<10}{'raw draft':<30}{'+glm editor':<30}")
    for side in ("flash", "pro"):
        raw, glm = matrix[f"{side}_raw"], matrix[f"{side}_glm"]
        rc = f"{raw['paired']} | {raw['marginal']}"
        gc = f"{glm['paired']} | {glm['marginal']}"
        print(f"  {side:<8}{rc:<30}{gc:<30}")
    d5 = result["d5_rule"]
    print(f"\n§D5 [{d5['status']}]: among COMPLETE traps flash-raw FAILS {d5['n_denominator']} "
          f"{d5['flash_fail_traps']} (min for verdict {d5['min_denominator']}); pro-raw FIXES {d5['pro_fixes']} "
          f"-> fix-rate {d5['fix_rate']} (thr 0.5); incomplete traps {d5['incomplete_traps']}")
    print(f"editor-delta pro  (glm-raw): {result['editor_delta_pro']['mean_delta']} "
          f"breaks={[b['id'] for b in result['editor_delta_pro']['editor_breaks']]}")
    print(f"editor-delta flash(glm-raw): {result['editor_delta_flash']['mean_delta']} "
          f"breaks={[b['id'] for b in result['editor_delta_flash']['editor_breaks']]}")
    print(f"seed-floor (DET non-unanimity rate): "
          f"{ {c: result['seed_floor'][c]['rate'] for c in CELLS} }")
    cov = result["coverage"]
    print(f"coverage (phenomenon RENDERED rate; omission proxy): "
          f"flash_raw {cov['flash_raw']['rate']} ({cov['flash_raw']['located']}/{cov['flash_raw']['seen']}) "
          f"pro_raw {cov['pro_raw']['rate']} ({cov['pro_raw']['located']}/{cov['pro_raw']['seen']}) | "
          f"flash_glm {cov['flash_glm']['rate']} pro_glm {cov['pro_glm']['rate']}")
    er = result["editor_restoration"]
    print(f"editor restoration (raw-omitted -> glm-rendered): "
          f"flash restored {er['flash']['n_restored']} broke {er['flash']['n_broke']} | "
          f"pro restored {er['pro']['n_restored']} broke {er['pro']['n_broke']}")
    print(f"COGS: flash-draft ${result['cogs'].get('draft_stage_flash')} vs "
          f"pro-draft ${result['cogs'].get('draft_stage_pro')} (x{result['cogs'].get('pro_over_flash_x')})")
    print("\nprovenance x-check (fresh flash_raw vs backend saved):")
    for t in traps:
        fresh = [per for per in per_trap if per["id"] == t["id"]][0]["cells"]["flash_raw"]["label"]
        bk = result["provenance_xcheck"]["backend_flash_raw"][t["id"]]
        armc = result["provenance_xcheck"]["exp14b_arm_C"][t["id"]]
        flag = "" if fresh == bk else "  <-- DIFF"
        print(f"  {t['id']:<3} fresh={fresh:<10} backend_flash={bk:<10} armC(flash+glm)={armc:<10}{flag}")


if __name__ == "__main__":
    main()
