#!/usr/bin/env python3
"""exp15 — pairwise span-citing trap judge (design §D3 / §1.6). Built from scratch: the exp14 harness
has NEITHER both-orders NOR 6->10 escalation.

Per trap-cell (one trap × one A0-vs-arm contrast), each of >=2 cross-family judges (grok-4.3 reasoning-ON
+ mistral-large-latest) is asked, BLIND (V1/V2), in BOTH orders, REP times, whether each translation
correctly resolves the SPECIFIC cohesion phenomenon (cite a <=15-word span). We record per-arm verdicts
(correct/partial/wrong -> accuracy) AND the pairwise winner AND catastrophe flags. Aggregation:
  base 6 votes/judge (3 reps × 2 orders); if a judge's decisive split is <=4:2 -> +4 to 10 (pre-reg rule,
  budget-relaxed from §A.7 "~10", fixed explicitly). Leave-one-judge-out. Catastrophe screen to the winner too.
Primary Q1/Q3 metric = per-arm ACCURACY (correct=1/partial=0.5/wrong=0), paired diff arm-A0. Pairwise
"better" is a robustness cross-check.

temp 0 (grok reasoning-ON is mildly stochastic across reps; mistral deterministic -> reps agree = strong
consensus, not wasted signal). Every call persists usage+cost via exp15_llm.Spender; per-call gate on.
"""
from __future__ import annotations
import json
import re
import sys
from dataclasses import dataclass, field
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
import exp15_llm as L

JUDGES = ["grok-4.3", "mistral-large-latest"]
BASE_REPS_PER_ORDER = 3          # 3 reps × 2 orders = 6 base votes/judge
ESCALATE_REPS_PER_ORDER = 2      # +2 reps × 2 orders = +4 -> 10 on split
SCORE = {"correct": 1.0, "partial": 0.5, "wrong": 0.0}

JUDGE_SYS = (
    "Ты — придирчивый билингвальный (китайский↔русский) литературный рецензент. Тебе дают ОДИН конкретный "
    "переводческий феномен связности (анафора, время/род глагола, катафора, эллипсис, перенос сущности через "
    "границу и т.п.), фрагмент КИТАЙСКОГО исходника и ДВА перевода этого места — V1 и V2. Оцени СТРОГО и ТОЛЬКО "
    "заданный феномен (не общий стиль). Для каждого перевода реши: correct (феномен передан верно), partial "
    "(частично/неоднозначно), wrong (искажён/потерян/неверный род-время-референт). Затем реши, какой перевод "
    "лучше справился ИМЕННО с этим феноменом: V1, V2 или tie. Отметь catastrophe=true, если перевод даёт грубую "
    "смысловую ошибку по этому феномену (неверный род/референт/полярность/потеря). Приведи доказательство — "
    "цитату <=15 слов из каждого перевода. Выведи СТРОГО JSON без пояснений:\n"
    '{"v1":"correct|partial|wrong","v2":"correct|partial|wrong","better":"V1|V2|tie",'
    '"cat_v1":true|false,"cat_v2":true|false,"ev_v1":"...","ev_v2":"..."}'
)


def _user_prompt(trap, v1_text, v2_text):
    return (
        f"ФЕНОМЕН СВЯЗНОСТИ ({trap['type']}): {trap['phenomenon']}\n"
        f"ЧТО ТРЕБУЕТСЯ ДЛЯ ВЕРНОЙ ПЕРЕДАЧИ: {trap['expected']}\n\n"
        f"КИТАЙСКИЙ ИСХОДНИК (место феномена):\n{trap['src_zone']}\n\n"
        f"V1 (перевод этого места):\n{v1_text}\n\n"
        f"V2 (перевод этого места):\n{v2_text}\n\n"
        "Оцени ТОЛЬКО указанный феномен. Ответ — строго JSON по схеме."
    )


_JSON_RE = re.compile(r"\{.*\}", re.S)


def _as_bool(x):
    """Robust truthiness for the catastrophe flags. A judge that emits the boolean as a JSON STRING
    ("cat_v1":"false") would otherwise become bool("false")==True -> a spurious catastrophe (review I4)."""
    return x is True or (isinstance(x, str) and x.strip().lower() in ("true", "1", "yes"))


def parse_verdict(text):
    m = _JSON_RE.search(text or "")
    if not m:
        return None
    try:
        d = json.loads(m.group(0))
    except Exception:
        # tolerant: pull ALL fields by regex (incl cat flags + evidence — review I4: the old fallback
        # dropped cat/ev, silently zeroing a flagged catastrophe on a parse-degraded but non-empty vote)
        d = {}
        for k in ("v1", "v2", "better"):
            mm = re.search(rf'"{k}"\s*:\s*"(\w+)"', text)
            if mm:
                d[k] = mm.group(1)
        for k in ("cat_v1", "cat_v2"):
            mm = re.search(rf'"{k}"\s*:\s*(true|false|"true"|"false")', text, re.I)
            if mm:
                d[k] = mm.group(1)
        for k in ("ev_v1", "ev_v2"):
            mm = re.search(rf'"{k}"\s*:\s*"([^"]*)"', text)
            if mm:
                d[k] = mm.group(1)
    for k in ("v1", "v2"):
        if d.get(k) not in SCORE:
            return None
    if d.get("better") not in ("V1", "V2", "tie"):
        d["better"] = "tie"
    d["cat_v1"] = _as_bool(d.get("cat_v1", False))
    d["cat_v2"] = _as_bool(d.get("cat_v2", False))
    return d


@dataclass
class Vote:
    judge: str
    order: str          # "AB" (a0=V1, arm=V2) or "BA" (arm=V1, a0=V2)
    a0_verdict: str     # correct/partial/wrong
    arm_verdict: str
    winner: str         # "a0" | "arm" | "tie"
    a0_cat: bool
    arm_cat: bool
    ev_a0: str = ""     # judge's <=15-word evidence span for the a0 side (order-normalised)
    ev_arm: str = ""    # judge's evidence span for the arm side
    raw: dict = field(default_factory=dict)


# a judge verdict is a tiny JSON (~60-150 completion tokens); grok reasoning is separate/uncapped by
# max_tokens. Cap max_tokens LOW so predict_cost's (max_out + reasoning_ceil) gate estimate stays under
# the judge per-call cap — otherwise grok's 8000-default made predict ~$0.033 > $0.03 and grok was
# silently gate-blocked (0 grok votes; the bug this selftest caught).
JUDGE_MAX_TOKENS = 1500


def _one_call(sp: L.Spender, judge, trap, v1_text, v2_text):
    """Return (verdict|None, rec, status). status in {ok, gate_block, err, parse_fail} so the caller can
    tell an out-of-budget refusal (gate_block) from a model parse failure (review I8)."""
    msgs = [{"role": "system", "content": JUDGE_SYS},
            {"role": "user", "content": _user_prompt(trap, v1_text, v2_text)}]
    in_est = int((len(JUDGE_SYS) + len(msgs[1]["content"])) / 1.9)   # conservative chars/token (review NIT)
    ok, pred, why = sp.gate(judge, in_est, JUDGE_MAX_TOKENS)
    if not ok:
        return None, {"judge": judge, "trap": trap["id"], "err": why, "cost": 0.0}, "gate_block"
    r = L.call(judge, msgs, temp=0.0, max_tokens=JUDGE_MAX_TOKENS)
    rec = {"judge": judge, "trap": trap["id"], "usage": r.get("usage"), "cost": r.get("cost", 0.0),
           "err": r.get("err"), "finish": r.get("finish"), "latency_s": r.get("latency_s"),
           "retry_fails": r.get("retry_fails")}
    sp.record(rec)
    import os
    if r.get("err"):
        if os.environ.get("JUDGE_DEBUG"):
            print(f"    [DBG {judge}] ERR {r['err'][:70]}")
        return None, rec, "err"
    v = parse_verdict(r["text"])
    if v is None:
        if os.environ.get("JUDGE_DEBUG"):
            print(f"    [DBG {judge}] NOPARSE content[{len(r.get('text') or '')}]={repr((r.get('text') or '')[:80])} "
                  f"finish={r.get('finish')} reasoning[{len(r.get('reasoning') or '')}]")
        return None, rec, "parse_fail"
    return v, rec, "ok"


def judge_cell(sp: L.Spender, trap, a0_text, arm_text, judges=JUDGES,
               base_reps=BASE_REPS_PER_ORDER, esc_reps=ESCALATE_REPS_PER_ORDER):
    """Return per-cell aggregate with per-arm accuracy, pairwise, catastrophe, per-judge (for LOO)."""
    votes: list[Vote] = []
    parse_fail = 0
    budget_truncated = False
    for judge in judges:
        if budget_truncated:
            break
        jvotes = []
        def do_reps(reps):
            nonlocal parse_fail, budget_truncated
            for _ in range(reps):
                if budget_truncated:
                    return
                for order in ("AB", "BA"):
                    if order == "AB":
                        v1, v2 = a0_text, arm_text
                    else:
                        v1, v2 = arm_text, a0_text
                    d, rec, status = _one_call(sp, judge, trap, v1, v2)
                    if status == "gate_block":
                        budget_truncated = True   # out of budget -> stop this cell, do NOT score partial
                        return
                    if d is None:
                        parse_fail += 1
                        continue
                    if order == "AB":
                        a0v, armv = d["v1"], d["v2"]
                        win = {"V1": "a0", "V2": "arm", "tie": "tie"}[d["better"]]
                        a0c, armc = bool(d["cat_v1"]), bool(d["cat_v2"])
                        eva0, evarm = d.get("ev_v1", ""), d.get("ev_v2", "")
                    else:
                        a0v, armv = d["v2"], d["v1"]
                        win = {"V1": "arm", "V2": "a0", "tie": "tie"}[d["better"]]
                        a0c, armc = bool(d["cat_v2"]), bool(d["cat_v1"])
                        eva0, evarm = d.get("ev_v2", ""), d.get("ev_v1", "")
                    jvotes.append(Vote(judge, order, a0v, armv, win, a0c, armc, eva0, evarm, d))
        do_reps(base_reps)
        # escalation: if this judge's decisive split is <=4:2 (winner has <=4 of 6 decisive), add reps
        dec = [v for v in jvotes if v.winner in ("a0", "arm")]
        arm_w = sum(1 for v in dec if v.winner == "arm")
        a0_w = sum(1 for v in dec if v.winner == "a0")
        if not budget_truncated and max(arm_w, a0_w) <= 4:   # split <=4:2 (or ties dominate) -> escalate
            do_reps(esc_reps)
        votes.extend(jvotes)
    return _aggregate(trap, votes, parse_fail, judges, budget_truncated)


def _acc(verdicts):
    return sum(SCORE[v] for v in verdicts) / len(verdicts) if verdicts else None


def _aggregate(trap, votes, parse_fail, judges, budget_truncated=False):
    a0_acc = _acc([v.a0_verdict for v in votes])
    arm_acc = _acc([v.arm_verdict for v in votes])
    arm_w = sum(1 for v in votes if v.winner == "arm")
    a0_w = sum(1 for v in votes if v.winner == "a0")
    tie = sum(1 for v in votes if v.winner == "tie")
    # leave-one-judge-out on the paired accuracy diff
    loo = {}
    for jx in judges:
        sub = [v for v in votes if v.judge != jx]
        if sub:
            loo[f"without_{jx}"] = round((_acc([v.arm_verdict for v in sub]) or 0)
                                         - (_acc([v.a0_verdict for v in sub]) or 0), 3)
    # order-bias check: winner distribution per order
    order_bias = {}
    for o in ("AB", "BA"):
        ov = [v for v in votes if v.order == o]
        order_bias[o] = {"arm": sum(1 for v in ov if v.winner == "arm"),
                         "a0": sum(1 for v in ov if v.winner == "a0"),
                         "tie": sum(1 for v in ov if v.winner == "tie")}
    return {
        "trap": trap["id"], "type": trap["type"],
        "n_votes": len(votes), "parse_fail": parse_fail, "budget_truncated": budget_truncated,
        "a0_accuracy": round(a0_acc, 3) if a0_acc is not None else None,
        "arm_accuracy": round(arm_acc, 3) if arm_acc is not None else None,
        "paired_diff_arm_minus_a0": round((arm_acc or 0) - (a0_acc or 0), 3) if a0_acc is not None else None,
        "pairwise": {"arm_wins": arm_w, "a0_wins": a0_w, "tie": tie},
        "catastrophe": {"a0": sum(1 for v in votes if v.a0_cat), "arm": sum(1 for v in votes if v.arm_cat)},
        "leave_one_judge_out": loo,
        "order_bias": order_bias,
        "per_judge": {j: {"arm_acc": round(_acc([v.arm_verdict for v in votes if v.judge == j]) or 0, 3),
                          "a0_acc": round(_acc([v.a0_verdict for v in votes if v.judge == j]) or 0, 3),
                          "n": sum(1 for v in votes if v.judge == j)} for j in judges},
        # per-vote raw verdicts (a0/arm-normalised, order recorded) — persisted so catastrophe/LOO/order
        # audits run on the actual votes, not just the aggregate counts (orchestrator D30.10 fix).
        "votes": [{"judge": v.judge, "order": v.order, "a0": v.a0_verdict, "arm": v.arm_verdict,
                   "winner": v.winner, "a0_cat": v.a0_cat, "arm_cat": v.arm_cat,
                   "ev_a0": v.ev_a0, "ev_arm": v.ev_arm} for v in votes],
    }


if __name__ == "__main__":
    print("judge harness module — import judge_cell(). Run selftest via judge_selftest.py.")
