#!/usr/bin/env python3
"""exp15/Q4a — SECONDARY grok+mistral judge mini-pass (design-compliance + judge-floor calibration
ON the meaning-traps). The $0 DET scorer (q4a_score.py) is PRIMARY and floor-immune; this pass exists
to (1) honour the frozen judge design (grok-4.3 reasoning-ON + mistral-large-latest, both orders,
6->10, LOO, catastrophe, per-vote persist) and (2) MEASURE how well judges even see these traps —
does the judge noise-floor (which swamped every S2' block at 0.126) also swamp DET-obvious meaning
errors? Reuses judges.judge_cell + exp15_llm verbatim (NO edits to those rigs).

Contrasts (--contrast):
  main   : a0=flash-raw(s0), arm=pro-raw(s0)     judge's pro-vs-flash draft signal (compare to DET)
  floor  : a0=flash-raw(s0), arm=flash-raw(s1)   two seeds of the SAME model = judge-noise + same-model
                                                 seed-content-variance UPPER BOUND (review F11): it is
                                                 pure judge noise ONLY where the two seeds DET-agree
                                                 (cross-referenced against q4a_score seed_floor).
  editor : a0=pro-raw(s0),   arm=pro-glm(s0)     judge's editor-delta (does editor break the draft)

Judge inputs are WINDOWED around the trap's phenomenon sentence (+/-700 chars, symmetric — flash and
pro translate the SAME source so the sentence sits at ~the same position; addendum §4 authorises
windowing for same-chunk symmetric contrasts) with a full-text fallback when the sentence isn't found.
This keeps the phenomenon visible while holding per-cell cost far under the per-call cap (a full ~5k
draft would push predict near the cap and could truncate the run mid-way).

Budget: OWN Spender, ledger q4a_judge_costs.jsonl, seed_total = current Q4a GEN spend (so the hard cap
is the GLOBAL $2 Q4a ceiling across gen+judge, NOT the exp15 $14.5). budget_truncated cells are DROPPED
from the analysis (review F3): a cell whose votes were cut off by the cap is partial/asymmetric and must
not enter the summary. per_call_cap kept high enough that ONLY the hard-cap governs (review F8). Per-call
gate + every completed vote persisted (judge gate-blocks are $0 and unledgered in the reused judges.py —
review F9, documented; the cutoff is recorded via budget_truncated in the results JSON)."""
from __future__ import annotations
import argparse
import json
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
import exp15_llm as L
import q4a_traps as Q
from judges import judge_cell

SD = Path("/home/ubuntu/books/gu-zhenren/exp15")
OUT = SD / "q4a"
GEN_LEDGER = SD / "q4a_gen_costs.jsonl"
JUDGE_LEDGER = SD / "q4a_judge_costs.jsonl"
Q4A_CAP = 2.0                    # GLOBAL Q4a ceiling ($; own budget, OUTSIDE exp15 $15 — orchestrator)

SEEDS = [0, 1, 2]
CONTRAST = {  # name -> (a0 cell, arm cell); the SEED is chosen per-trap (see _pick_seed)
    "main":   ("flash_raw", "pro_raw"),      # judge's pro-vs-flash draft signal
    "floor":  ("flash_raw", "flash_raw"),    # same model, two DIFFERENT seeds -> noise+seed-variance floor
    "editor": ("pro_raw", "pro_glm"),        # judge's editor-delta
}


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 _located(cell, cid, seed, trap):
    txt = _read(cell, cid, seed)
    return bool(txt) and bool(Q.sentence_with(txt, trap["locator"]))


def _pick_seed(a0cell, armcell, trap, floor):
    """Per-trap seed selection: judge the meaning GIVEN both sides RENDERED the phenomenon (omission is a
    separate DET-coverage metric). Return (a0seed, armseed) or None. For 'floor', a0/arm are the same cell
    -> pick two DISTINCT seeds both rendering the trap (else the floor conflates omission with judge noise)."""
    cid = trap["chunk"]
    if floor:
        rs = [s for s in SEEDS if _located(a0cell, cid, s, trap)]
        return (rs[0], rs[1]) if len(rs) >= 2 else None
    for s in SEEDS:                                  # same seed both sides, both rendering
        if _located(a0cell, cid, s, trap) and _located(armcell, cid, s, trap):
            return (s, s)
    return None


def _window(text, trap, half=700):
    """Symmetric +/-half window around the trap's phenomenon sentence; full text if not located."""
    sent = Q.sentence_with(text, trap["locator"])
    if not sent:
        return text
    probe = sent[:40]
    i = text.find(probe)
    if i < 0:
        return text
    lo, hi = max(0, i - half), min(len(text), i + len(sent) + half)
    return text[lo:hi]


def _gen_spent():
    t = 0.0
    if GEN_LEDGER.exists():
        for line in GEN_LEDGER.read_text(encoding="utf-8").splitlines():
            try:
                t += json.loads(line).get("cost", 0.0)
            except Exception:
                pass
    return t


def run(contrast, per_call_cap, limit=0):
    traps = Q.traps()
    if limit:
        traps = traps[:limit]
    a0cell, armcell = CONTRAST[contrast]
    gen = _gen_spent()
    sp = L.Spender(str(JUDGE_LEDGER), per_call_cap=per_call_cap, hard_cap=Q4A_CAP, seed_total=gen)
    print(f"[budget] Q4a gen spent ${gen:.4f}; judge ledger start ${sp.total:.4f}; global cap ${Q4A_CAP}; "
          f"contrast={contrast} (a0={a0cell} arm={armcell}, seed chosen per-trap where both render)")
    out, skipped, truncated = [], [], []
    for t in traps:
        pick = _pick_seed(a0cell, armcell, t, floor=(contrast == "floor"))
        if pick is None:
            skipped.append((t["id"], "no-seed-both-render")); continue
        a0seed, armseed = pick
        a0_full = _read(a0cell, t["chunk"], a0seed)
        arm_full = _read(armcell, t["chunk"], armseed)
        det_a0, _ = Q.det_score(a0_full, t)
        det_arm, _ = Q.det_score(arm_full, t)
        agg = judge_cell(sp, Q.judge_trap(t), a0_text=_window(a0_full, t), arm_text=_window(arm_full, t))
        if agg.get("budget_truncated"):
            # DROP a cap-truncated (partial/asymmetric) cell — do NOT let it enter the analysis (review F3)
            truncated.append(t["id"])
            print(f"  [HARD-CAP] Q4a ${Q4A_CAP} bound at {t['id']} — cell dropped, stopping."); break
        agg.update(contrast=contrast, id=t["id"], cat_class=t["catastrophe_class"],
                   a0_seed=a0seed, arm_seed=armseed, det_a0=det_a0, det_arm=det_arm)
        out.append(agg)
        print(f"  {contrast} {t['id']} (s{a0seed}/s{armseed}): judge diff(arm-a0) {agg['paired_diff_arm_minus_a0']} "
              f"(a0_acc {agg['a0_accuracy']} arm_acc {agg['arm_accuracy']}; DET a0={det_a0} arm={det_arm})"
              f"  spend ${sp.total:.4f}")
    resfile = SD / f"q4a_judge_results_{contrast}.json"
    resfile.write_text(json.dumps({"contrast": contrast, "n": len(out), "skipped": skipped,
                                   "truncated_dropped": truncated, "results": out},
                                  ensure_ascii=False, indent=2), encoding="utf-8")
    # quick summary
    decided = [a for a in out if a["paired_diff_arm_minus_a0"] is not None]
    if decided:
        import statistics
        diffs = [a["paired_diff_arm_minus_a0"] for a in decided]
        print(f"\n{contrast}: n={len(decided)} mean diff {round(statistics.mean(diffs),3)} "
              f"mean|diff| {round(statistics.mean(abs(d) for d in diffs),3)}  spend ${sp.total:.4f}")
    print(f"{contrast}: {len(out)} scored, {len(skipped)} skipped {skipped}, "
          f"{len(truncated)} cap-dropped {truncated}")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--contrast", choices=list(CONTRAST), required=True)
    # per-call cap high enough that ONLY the $2 hard-cap governs (windowed inputs predict ~$0.03; review F8)
    ap.add_argument("--per-call-cap", type=float, default=0.09)
    ap.add_argument("--limit", type=int, default=0)
    a = ap.parse_args()
    run(a.contrast, a.per_call_cap, a.limit)


if __name__ == "__main__":
    main()