#!/usr/bin/env python3
"""Reconcile the code judge against the trusted Opus judge.

fox 2026-05-19: "we do trust opus from claude highly even though it
costs a lot." The B rescore (control_sweep_*_code_judge.jsonl) carries
both the original Opus verdict (`verdict` field, from the live sweep)
and the deterministic code verdict (`code_verdict` field, from the
rescore). Where the two disagree, Opus is the reference — the code
judge is the instrument being calibrated.

This script:

  1. bucketises disagreements by (code_verdict, opus_verdict) cell;
  2. dumps a configurable sample from each cell with full context
     (question, answer, code rationale, opus rationale-tail);
  3. emits a markdown report so the patterns are reviewable
     side-by-side rather than scrolling per-record.

Reads only records where Opus produced a real verdict
(verdict != 'JUDGE_ERROR') — the rest is Opus failing, not
disagreement.

NO LLM calls. Pure read of the JSONL.

Usage::

    python -m bench.analyze_judge_disagreement \\
        --in bench/qa_results/control_sweep_<TS>_code_judge.jsonl \\
        --sample 10
"""
from __future__ import annotations

import argparse
import json
import sys
from collections import Counter, defaultdict
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
from bench.control_ab import _gold  # noqa: E402

LABELS = ["CORRECT_GROUNDED", "WRONG", "FABRICATED", "ABSTAINED",
          "JUDGE_ERROR"]


def _opus_rationale_tail(judge_raw: str) -> str:
    """Heuristic: the last non-sentinel non-empty line of the Opus raw
    output is usually the model's brief rationale before the
    FINAL_VERDICT=X sentinel."""
    out = []
    for ln in reversed((judge_raw or "").splitlines()):
        s = ln.strip()
        if not s:
            continue
        if "FINAL_VERDICT" in s:
            continue
        out.append(s)
        if len(out) >= 2:
            break
    return " · ".join(reversed(out))[:300]


def _gold_excerpt(shards_dir: Path, fixture_by_i: dict, i: int,
                  cap: int = 600, cache: dict | None = None) -> str:
    cache = cache if cache is not None else {}
    if i in cache:
        return cache[i]
    fxr = fixture_by_i.get(i)
    if not fxr:
        cache[i] = ""
        return ""
    g = _gold(shards_dir, fxr.get("shard", ""), fxr["target_root"]) or ""
    cache[i] = g[:cap]
    return cache[i]


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--in", dest="in_path", required=True)
    ap.add_argument("--fixture",
                    default="bench/qa_questions_stale_map.json")
    ap.add_argument("--shards-dir",
                    default=str(Path.home() / ".arborist" / "shards"))
    ap.add_argument("--sample", type=int, default=8,
                    help="per-cell sample size to dump")
    ap.add_argument("--out-md", default="",
                    help="markdown report path "
                         "(default: alongside --in with _disagreement.md)")
    a = ap.parse_args()

    in_path = Path(a.in_path)
    out_md = (Path(a.out_md) if a.out_md
              else in_path.with_name(
                  in_path.stem.replace("_code_judge", "")
                  + "_judge_disagreement.md"))

    items = json.loads(Path(a.fixture).read_text())
    fixture_by_i = {i + 1: it for i, it in enumerate(items)}
    shards_dir = Path(a.shards_dir)
    gold_cache: dict[int, str] = {}

    cells: dict[tuple, list[dict]] = defaultdict(list)
    cell_count: Counter = Counter()
    opus_total: Counter = Counter()

    for ln in in_path.read_text().splitlines():
        if not ln.strip():
            continue
        r = json.loads(ln)
        if r.get("arm") == "skip":
            continue
        opus = r.get("verdict") or "JUDGE_ERROR"
        code = r.get("code_verdict") or "JUDGE_ERROR"
        opus_total[opus] += 1
        # We're calibrating against records Opus actually graded.
        if opus == "JUDGE_ERROR":
            continue
        cell_count[(code, opus)] += 1
        cells[(code, opus)].append(r)

    # Bucketise into "agreement" (on-diagonal), "off-diagonal", and
    # rank off-diagonal cells by count so the biggest disagreements
    # surface first.
    on_diag = sum(cell_count[(lab, lab)] for lab in LABELS)
    total = sum(cell_count.values())
    off_diag_cells = sorted(
        ((cnt, code, opus) for (code, opus), cnt in cell_count.items()
         if code != opus),
        reverse=True,
    )

    L: list[str] = [
        f"# Judge disagreement reconciliation — `{in_path.name}`",
        "",
        "fox 2026-05-19: Opus is trusted; code judge is being calibrated. "
        "Below: every (code_verdict, opus_verdict) cell off the diagonal, "
        f"sample={a.sample} records per cell, sorted by cell size "
        "(biggest disagreement first).",
        "",
        f"**Restricted to records Opus actually graded** "
        f"(`verdict != 'JUDGE_ERROR'`): {total} of "
        f"{sum(opus_total.values())} (the other "
        f"{opus_total['JUDGE_ERROR']} are Opus failures, not "
        "disagreement).",
        "",
        f"On-diagonal (agree): **{on_diag}/{total} = "
        f"{on_diag/total:.1%}**.",
        "",
        "## Per-cell summary",
        "",
        "| n | code says | opus says |",
        "|---|---|---|"
    ]
    for cnt, code, opus in off_diag_cells:
        L.append(f"| {cnt} | {code} | {opus} |")
    L += [""]

    # Per-cell dumps, biggest first.
    for cnt, code, opus in off_diag_cells:
        if cnt == 0:
            continue
        L += [
            f"## code=**{code}** · opus=**{opus}** "
            f"({cnt} records · showing {min(a.sample, cnt)})",
            "",
        ]
        for r in cells[(code, opus)][:a.sample]:
            i = r.get("i")
            q = r.get("question_asked") or r.get("question_orig") or ""
            ans = (r.get("answer") or "").replace("\n", " ")[:300]
            cr = (r.get("code_rationale") or "")[:200]
            jr = _opus_rationale_tail(r.get("judge_raw") or "")[:200]
            gold = _gold_excerpt(shards_dir, fixture_by_i, i,
                                 cap=400, cache=gold_cache)
            gold_short = gold.replace("\n", " ")[:300]
            L += [
                f"### item {i} · {r.get('arm')}/{r.get('model')}/{r.get('variant')}",
                f"- **Q:** {q}",
                f"- **A:** {ans}",
                f"- **Gold excerpt:** {gold_short}",
                f"- **Code says {code}:** {cr}",
                f"- **Opus says {opus}:** {jr}",
                "",
            ]

    out_md.parent.mkdir(parents=True, exist_ok=True)
    out_md.write_text("\n".join(L) + "\n")
    print(f"  disagreement report: {out_md}")
    print(f"  on-diagonal: {on_diag}/{total} = {on_diag/total:.1%}")
    for cnt, code, opus in off_diag_cells[:10]:
        print(f"    n={cnt:>4}  code={code:<20} opus={opus}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
