#!/usr/bin/env python3
"""exp15 — judge driver for the carrying block (Q1b oracle-vs-flat + Q3a carryover + Q0 + FLOOR), §D3.

Maps each final_traps.json trap to the chunk whose SOURCE contains its dependent_span, in each of:
  flat-A0-30 (blind baseline)   oracle-30 (seam-aligned)   arm modes q3a-src/draft/final
then gives the pairwise span judge (judges.py: both-orders, 6->10, LOO, catastrophe, per-vote persist)
the trap + a DEPENDENT-CENTRED SYMMETRIC WINDOW of each side's translation.

WINDOW POLICY (orchestrator CRITICAL fix, replaces the old HEAD_CHARS=1100 head-excerpt):
  The old head-excerpt showed the judge only the first 1100 chars of each final. In the FLAT arm the
  straddle-dependent lands at frac ~0 (chunk head) -> always visible. In the ORACLE arm the SAME
  dependent sits mid-chunk (frac 0.09-1.0) -> its Russian rendering fell PAST char 1100 in 13/18
  T-ana cells, so the judge scored "absent"=wrong/catastrophe by construction. That artifact, not the
  boundary, produced the -0.564 headline. Fix: a fixed-length window centred on the dependent's
  PROPORTIONAL position in each side's target, SAME length L for every arm (size-symmetric). Both
  sides now always contain the local phenomenon (anchor+dependent sit within a few hundred source
  chars of each other -> inside one window). The floor block (A0 vs A0') validates the policy: if the
  window systematically hid phenomena the floor would be large/asymmetric.

Q1b:   a0=flat, arm=oracle      (does cutting AT the scene boundary help vs a blind mid-scene cut?)
Q3a:   a0=flat, arm=mode        (does carryover fix the cross-boundary dependency? per mode)
Q0 :   a0=do-nothing (draft), arm=A0 (editor final)   (TOST equivalence — APE lesson)
FLOOR: a0=flat-A0, arm=flat-A0' (independent regen, same config) — the judge+generation NOISE floor;
       effects below it are not interpretable (freeze §1.5). Same window policy as every block.

Oracle mapping is REQUIRED only for Q1b. Flat mapping is mandatory for all blocks. (Old code required
BOTH -> T-ana-33, whose oracle chunk 12 was cjk_artifact-skipped, was silently dropped from Q3a/Q0/FLOOR
too. Now flat-only blocks keep it; Q1b documents its exclusion.)

Judge SUBSET (budget scope, guard 5): ALL T-ana (primary) + ALL T-ell + a capped T-ent CONTROL sample.
  --project : $0 cell/token/cost estimate (guard-5 projection; NO api calls)
  --run Q1b|Q3a|Q0|FLOOR : execute that block   |   --tent-control N : cap the T-ent control cells
"""
from __future__ import annotations
import argparse
import json
import sqlite3
import sys
from pathlib import Path

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

SD = Path("/home/ubuntu/books/gu-zhenren/exp15")
AH = SD / "anchors"
ARMS = SD / "arms"
WIN_CHARS = 1500           # dependent-centred symmetric window fed to the judge (same L for both arms)
GLOBAL_CEILING = 14.5      # hard global spend ceiling across ALL blocks/ledgers (review I8; $15 freeze - margin)
# Orchestrator's AUTHORITATIVE pre-session-2 spend (2541 billing lines, 0 discrepancies). Used as the seed
# floor because it includes the tmctl anchor-generation (flat/oracle/a0prime .db) that ran through the Go
# backend, NOT the Python Spender ledgers — so a raw ledger-sum would under-count the true global by ~$2.
ALREADY_SPENT = 9.82
NEW_LEDGERS = ("judge_costs_FLOOR.jsonl", "judge_costs_Q1b.jsonl", "judge_costs_Q4a.jsonl")


def global_spent_excluding(block):
    """ALREADY_SPENT floor + this session's NEW re-run ledgers (except the current block's own, which its
    own Spender loads). Enforces a true global ceiling via seed_total (review I8)."""
    tot = ALREADY_SPENT
    for name in NEW_LEDGERS:
        if name == f"judge_costs_{block}.jsonl":
            continue
        p = SD / name
        if p.exists():
            for line in p.read_text(encoding="utf-8").splitlines():
                try:
                    tot += json.loads(line).get("cost", 0.0)
                except Exception:
                    pass
    return tot


def _db_stage(db, stage):
    con = sqlite3.connect(db)
    out = {}
    for ck, txt in con.execute("SELECT chunk_idx, response_text FROM checkpoints WHERE stage=? "
                               "ORDER BY chunk_idx, attempt", (stage,)):
        out[ck] = txt
    con.close()
    return out


def db_finals(db):
    return _db_stage(db, "edit")


def db_drafts(db):
    return _db_stage(db, "draft")


def find_off(chunks, span):
    """(chunk_idx, char_offset) of the first chunk whose source contains span's 20-char head."""
    probe = (span or "").strip()[:20]
    for i, c in enumerate(chunks):
        if probe and probe in c:
            return i, c.find(probe)
    return None, None


def locate_flat(trap, flat_src):
    """Prefer the trap's authoritative dep_chunk (verified to contain the dependent); else scan."""
    d = (trap["dependent_span"] or "").strip()[:20]
    fi = trap.get("dep_chunk")
    if fi is not None and 0 <= fi < len(flat_src) and d and d in flat_src[fi]:
        return fi, flat_src[fi].find(d)
    return find_off(flat_src, trap["dependent_span"])


def frac(chunk_src, off):
    n = len(chunk_src or "")
    return (off / n) if (n and off is not None and off >= 0) else 0.0


def window(text, fr, L=WIN_CHARS):
    """Fixed-length (L) window of `text` centred on proportional position `fr`, clamped-and-extended
    so it is always exactly L chars when text is longer than L (size-symmetric across arms)."""
    text = text or ""
    if len(text) <= L:
        return text
    center = int(round(fr * len(text)))
    lo, hi = center - L // 2, center - L // 2 + L
    if lo < 0:
        lo, hi = 0, L
    if hi > len(text):
        hi, lo = len(text), len(text) - L
    return text[lo:hi]


def load():
    flat_src = greedy_chunks((SD / "S2prime_big_flat.txt").read_text(encoding="utf-8"))
    oracle_src = greedy_chunks((SD / "S2prime_big_blanklines.txt").read_text(encoding="utf-8"))
    flat_fin = db_finals(AH / "flat_a0_30.db")
    oracle_fin = db_finals(AH / "oracle_30.db")
    traps = json.load(open(SD / "final_traps.json"))["traps"]
    return flat_src, oracle_src, flat_fin, oracle_fin, traps


def arm_final(mode, chunk_idx):
    p = ARMS / mode / f"{chunk_idx:02d}.final.txt"
    return p.read_text(encoding="utf-8") if p.exists() else None


def locate_oracle(trap, oracle_src):
    """Select the oracle chunk holding the phenomenon. Prefer a chunk containing BOTH the anchor and the
    dependent (seam-aligned -> co-located for T-ana/T-ell). Else accept a dependent-only chunk ONLY when
    the dependent probe is distinctive (>=8 chars) AND unique — short/ambiguous probes are skipped rather
    than judged on the wrong passage (review I3). Returns (chunk_idx, dep_offset) or (None, None)."""
    d = (trap["dependent_span"] or "").strip()[:20]
    a = (trap["anchor_span"] or "").strip()[:20]
    for i, c in enumerate(oracle_src):
        if d and a and d in c and a in c:
            return i, c.find(d)
    hits = [i for i, c in enumerate(oracle_src) if d and d in c]
    if len(d) >= 8 and len(hits) == 1:
        return hits[0], oracle_src[hits[0]].find(d)
    return None, None


def excerpt(final, fr, full):
    """Full final (no windowing) when full=True; else the dependent-centred window."""
    return final if full else window(final, fr)


def map_trap(trap, flat_src, oracle_src, flat_fin, oracle_fin, full=False):
    """Flat mapping is mandatory (returns None if absent). Oracle mapping is optional (None if absent
    -> Q1b skips, flat-only blocks proceed). full=True feeds whole finals (no windowing)."""
    fi, fpos = locate_flat(trap, flat_src)
    if fi is None or fi not in flat_fin:
        return None
    ffr = frac(flat_src[fi], fpos)
    jt = {"id": trap["id"], "type": trap["type"],
          "src_zone": trap["anchor_span"] + " … " + trap["dependent_span"],
          "phenomenon": trap["phenomenon"], "expected": trap["expected_ru"]}
    orac = None
    oi, opos = locate_oracle(trap, oracle_src)
    if oi is not None and oi in oracle_fin:
        ofr = frac(oracle_src[oi], opos)
        orac = {"chunk": oi, "frac": round(ofr, 3), "win": excerpt(oracle_fin[oi], ofr, full)}
    return {"jt": jt, "flat_chunk": fi, "flat_frac": round(ffr, 3),
            "flat_win": excerpt(flat_fin[fi], ffr, full), "oracle": orac}


def select_traps(traps, tent_control, primary_only=False):
    prim = [t for t in traps if t["type"] in ("T-ana", "T-ell")]
    if primary_only:
        return prim
    tent = [t for t in traps if t["type"] == "T-ent"][:tent_control]
    return prim + tent


def project(traps, tent_control):
    flat_src, oracle_src, flat_fin, oracle_fin, _ = load()
    sub = select_traps(traps, tent_control)
    mapped = [t for t in sub if map_trap(t, flat_src, oracle_src, flat_fin, oracle_fin)]
    n = len(mapped)
    q1b_ok = sum(1 for t in mapped
                 if (map_trap(t, flat_src, oracle_src, flat_fin, oracle_fin) or {}).get("oracle"))
    in_tok = (600 + 400 + 2 * WIN_CHARS // 2)            # sys + trap + 2 windows (~2 chars/token)
    votes = 7.4                                          # 6 base + 0.35*4 escalation avg, per judge x2
    g = (in_tok * L._price("grok-4.3", "in") + 900 * L._price("grok-4.3", "out")) / 1e6
    m = (in_tok * L._price("mistral-large-latest", "in") + 200 * L._price("mistral-large-latest", "out")) / 1e6
    cc = votes * (g + m)
    q3a_cells = sum(1 for t in mapped if t["type"] in ("T-ana", "T-ell")) * 3 + \
                sum(1 for t in mapped if t["type"] == "T-ent")
    print(f"=== guard-5 judge projection (WIN_CHARS={WIN_CHARS}, tent_control={tent_control}) ===")
    print(f"subset {len(sub)}  flat-mappable {n} (T-ana {sum(1 for t in mapped if t['type']=='T-ana')}, "
          f"T-ell {sum(1 for t in mapped if t['type']=='T-ell')}, T-ent {sum(1 for t in mapped if t['type']=='T-ent')})  "
          f"oracle-mappable {q1b_ok}")
    print(f"per-cell in_tok~{in_tok} votes~{votes} cost/cell ~${cc:.4f}")
    print(f"Q1b(re-judge) cells {q1b_ok} -> ~${q1b_ok*cc:.2f}")
    print(f"FLOOR cells {n} -> ~${n*cc:.2f}")
    print(f"Q3a cells {q3a_cells} -> ~${q3a_cells*cc:.2f}  |  Q0 cells {n} -> ~${n*cc:.2f}")


def run(block, tent_control, per_call_cap, limit=0, full=False):
    flat_src, oracle_src, flat_fin, oracle_fin, traps = load()
    primary_only = block in ("FLOOR",)
    sub = select_traps(traps, tent_control, primary_only=primary_only)
    if limit:
        sub = sub[:limit]
    a0p_fin = db_finals(AH / "flat_a0prime_30.db") if block == "FLOOR" else None
    draft_fin = db_drafts(AH / "flat_a0_30.db") if block == "Q0" else None
    # GLOBAL cap: seed this block's Spender with spend already booked in all OTHER ledgers so hard_cap is
    # a true cross-block ceiling (review I8), not a per-block one that would sum to ~4x the freeze.
    seed = global_spent_excluding(block)
    sp = L.Spender(str(SD / f"judge_costs_{block}.jsonl"), per_call_cap=per_call_cap,
                   hard_cap=GLOBAL_CEILING, seed_total=seed)
    print(f"[budget] global spent (other ledgers) ${seed:.3f}; block start ${sp.total:.3f}; "
          f"global ceiling ${GLOBAL_CEILING}; excerpt={'FULL-CHUNK' if full else f'window{WIN_CHARS}'}")
    out, skipped = [], []
    for t in sub:
        m = map_trap(t, flat_src, oracle_src, flat_fin, oracle_fin, full=full)
        if not m:
            skipped.append((t["id"], "flat-unmappable")); continue
        dep = m["flat_chunk"]
        if block == "Q1b":
            if not m["oracle"]:
                skipped.append((t["id"], "oracle-unmappable")); continue
            agg = judge_cell(sp, m["jt"], a0_text=m["flat_win"], arm_text=m["oracle"]["win"])
            agg.update(block="Q1b", arm="oracle", flat_chunk=dep, flat_frac=m["flat_frac"],
                       oracle_chunk=m["oracle"]["chunk"], oracle_frac=m["oracle"]["frac"], full=full)
            out.append(agg)
            print(f"  Q1b {t['id']}: oracle-flat {agg['paired_diff_arm_minus_a0']} "
                  f"(ffrac {m['flat_frac']} ofrac {m['oracle']['frac']} trunc={agg.get('budget_truncated')})  spend ${sp.total:.3f}")
        elif block == "FLOOR":
            if dep not in a0p_fin:
                skipped.append((t["id"], "a0prime-missing")); continue
            agg = judge_cell(sp, m["jt"], a0_text=m["flat_win"],
                             arm_text=excerpt(a0p_fin[dep], m["flat_frac"], full))
            agg.update(block="FLOOR", arm="a0prime", flat_chunk=dep, flat_frac=m["flat_frac"], full=full)
            out.append(agg)
            print(f"  FLOOR {t['id']}: |diff| {abs(agg['paired_diff_arm_minus_a0'] or 0):.3f}  spend ${sp.total:.3f}")
        elif block == "Q3a":
            modes = ["q3a-src", "q3a-draft", "q3a-final"] if t["type"] in ("T-ana", "T-ell") else ["q3a-final"]
            for mode in modes:
                af = arm_final(mode, dep)
                if not af:
                    skipped.append((t["id"], f"{mode}-missing")); continue
                agg = judge_cell(sp, m["jt"], a0_text=m["flat_win"], arm_text=excerpt(af, m["flat_frac"], full))
                agg.update(block="Q3a", arm=mode, flat_chunk=dep, flat_frac=m["flat_frac"], full=full)
                out.append(agg)
                print(f"  Q3a {t['id']} {mode}: {agg['paired_diff_arm_minus_a0']}  spend ${sp.total:.3f}")
        elif block == "Q0":
            if dep not in draft_fin:
                skipped.append((t["id"], "draft-missing")); continue
            agg = judge_cell(sp, m["jt"], a0_text=excerpt(draft_fin[dep], m["flat_frac"], full),
                             arm_text=m["flat_win"])   # a0=do-nothing(draft), arm=A0(editor final)
            agg.update(block="Q0", arm="A0-editor", flat_chunk=dep, flat_frac=m["flat_frac"], full=full)
            out.append(agg)
            print(f"  Q0 {t['id']}: A0-donothing {agg['paired_diff_arm_minus_a0']}  spend ${sp.total:.3f}")
        if any(a.get("budget_truncated") for a in out):
            print(f"  [HARD-CAP] budget ceiling ${GLOBAL_CEILING} bound mid-run -> stopping block"); break
    Path(SD / f"judge_results_{block}.json").write_text(
        json.dumps({"n": len(out), "win_chars": (None if full else WIN_CHARS), "full_chunk": full,
                    "skipped": skipped, "results": out}, ensure_ascii=False, indent=2), encoding="utf-8")
    print(f"\n{block}: {len(out)} cells judged, {len(skipped)} skipped {skipped}, block+global spend ${sp.total:.4f}")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--project", action="store_true")
    ap.add_argument("--run", choices=["Q1b", "Q3a", "Q0", "FLOOR"])
    ap.add_argument("--tent-control", type=int, default=10)
    ap.add_argument("--per-call-cap", type=float, default=0.03)
    ap.add_argument("--full", action="store_true", help="feed WHOLE finals (no windowing) — used for Q1b")
    ap.add_argument("--limit", type=int, default=0, help="judge only first N traps (selftest)")
    a = ap.parse_args()
    traps = json.load(open(SD / "final_traps.json"))["traps"]
    if a.project:
        project(traps, a.tent_control)
    elif a.run:
        run(a.run, a.tent_control, a.per_call_cap, a.limit, full=a.full)
    else:
        ap.print_help()


if __name__ == "__main__":
    main()
