#!/usr/bin/env python3
"""Re-grade an existing #000057 control-sweep JSONL with the code judge.

fox 2026-05-19: "we burned all our Opus … let's get data first and we
can use tools to judge before passing samples to Opus later." The sweep
JSONLs from the huge-N run carry every (question, answer) but NOT the
gold source (gold is re-fetched from shards at grade time so the JSONL
stays compact). This script:

  1. loads a sweep JSONL (the recorded Opus verdicts on each record);
  2. for each record, re-fetches gold via ``bench.control_ab._gold``
     using the matching fixture entry (matched by item index ``i``);
  3. runs ``bench.judge_code.judge`` on (question_asked, answer, gold);
  4. emits an agreement matrix: code-judge label × opus-judge label,
     plus a per-arm / per-model / per-variant scorecard;
  5. emits the JUDGE_ERROR residue list — the records the code judge
     could not classify deterministically. That residue is the eventual
     LLM-batch needle-haystack ask. If the residue is small, you do not
     need a paid second pass; if it's big, you know exactly the size of
     the spend you would be authorising.

NO LLM calls, no shard writes. Reads the existing JSONL & shards
read-only; writes a markdown scorecard + a JSONL of code-judge verdicts
(one row per existing-sweep record, joinable on ``i,arm,model,variant``).

Usage::

    python -m bench.score_with_code_judge \\
        --in bench/qa_results/control_sweep_<TS>.jsonl \\
        --fixture bench/qa_questions_stale_map.json \\
        --shards-dir ~/.arborist/shards \\
        --out-md bench/results/code_judge_rescore_<TS>.md \\
        --out-jsonl bench/qa_results/code_judge_rescore_<TS>.jsonl

The default fixture matches the sweep's default; if you ran the sweep
with a different ``--fixture`` you must pass the same one here.
"""
from __future__ import annotations

import argparse
import json
import sys
import time
from collections import Counter
from pathlib import Path

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


def _load_jsonl(p: Path) -> list[dict]:
    """Tolerant load — same convention as control_sweep._load_recs."""
    out = []
    for ln in p.read_text().splitlines():
        if not ln.strip():
            continue
        try:
            out.append(json.loads(ln))
        except json.JSONDecodeError:
            continue
    return out


def _agree_md(matrix: Counter, code_total: Counter, opus_total: Counter) -> list[str]:
    labels = ["CORRECT_GROUNDED", "WRONG", "FABRICATED", "ABSTAINED",
              "JUDGE_ERROR"]
    L = ["## Agreement matrix (rows = code judge, columns = original judge in sweep)",
         "",
         "| code \\ orig | " + " | ".join(labels) + " | row total |",
         "|" + "|".join(["---"] * (len(labels) + 2)) + "|"]
    for cl in labels:
        row = [str(matrix.get((cl, ol), 0)) for ol in labels]
        L.append(f"| **{cl}** | " + " | ".join(row) + f" | {code_total[cl]} |")
    L.append("| **col total** | " + " | ".join(
        str(opus_total[ol]) for ol in labels) + " | |")
    # Diagonal = exact agreement; off-diagonal = disagreement.
    diag = sum(matrix.get((lab, lab), 0) for lab in labels)
    total = sum(matrix.values())
    if total:
        L += ["",
              f"Exact-label agreement: **{diag}/{total} = "
              f"{diag/total:.1%}**."]
    return L


def _bucket_str(c: Counter) -> str:
    tot = sum(c.values()) or 1
    return (f"CG={c['CORRECT_GROUNDED']} W={c['WRONG']} "
            f"F={c['FABRICATED']} A={c['ABSTAINED']} "
            f"JE={c['JUDGE_ERROR']}  (n={tot})")


def _scorecard_md(by_key: dict, title: str) -> list[str]:
    L = [f"## {title}",
         "",
         "Each row is the same set of records as the matching column "
         "of the original sweep — only the judge changed.",
         "",
         "| key | code-judge verdicts |",
         "|---|---|"]
    for k, c in sorted(by_key.items()):
        L.append(f"| `{k}` | {_bucket_str(c)} |")
    return L


def _residue_md(residue: list[dict], cap: int = 30) -> list[str]:
    L = [f"## JUDGE_ERROR residue (top {min(cap, len(residue))} of "
         f"{len(residue)})",
         "",
         "These are the records the code judge could NOT classify "
         "deterministically. They are the natural input to a later "
         "LLM-batch needle-haystack pass (Opus or Grok). Each row "
         "shows the code-judge rationale + a short answer excerpt so "
         "the residue is greppable / sortable before any LLM spend.",
         "",
         "| i | arm | model | variant | original | code rationale | answer |",
         "|---|---|---|---|---|---|---|"]
    for r in residue[:cap]:
        ans = (r.get("answer") or "").replace("|", "/")[:80]
        L.append(f"| {r.get('i')} | {r.get('arm')} | {r.get('model')} | "
                 f"{r.get('variant')} | {r.get('verdict')} | "
                 f"{(r.get('code_rationale') or '')[:90]} | {ans} |")
    return L


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--in", dest="in_path", required=True,
                    help="sweep JSONL to re-grade")
    ap.add_argument("--fixture",
                    default="bench/qa_questions_stale_map.json",
                    help="must match the fixture the sweep was run on "
                         "(records carry item index i, not gold)")
    ap.add_argument("--shards-dir",
                    default=str(Path.home() / ".arborist" / "shards"))
    ap.add_argument("--out-md", default="",
                    help="markdown scorecard output path "
                         "(default: alongside --in with _code_judge.md)")
    ap.add_argument("--out-jsonl", default="",
                    help="JSONL of per-record code-judge verdicts "
                         "(default: alongside --in with _code_judge.jsonl)")
    a = ap.parse_args()

    in_path = Path(a.in_path)
    if not in_path.exists():
        print(f"ABORT: {in_path} not found", file=sys.stderr)
        return 1
    out_md = Path(a.out_md) if a.out_md else \
        in_path.with_name(in_path.stem + "_code_judge.md")
    out_jsonl = Path(a.out_jsonl) if a.out_jsonl else \
        in_path.with_name(in_path.stem + "_code_judge.jsonl")

    items = json.loads(Path(a.fixture).read_text())
    # Index `i` in the sweep records is 1-based (enumerate start=1) and
    # references items in fixture order; rebuild the lookup table.
    by_i = {i + 1: it for i, it in enumerate(items)}

    recs = _load_jsonl(in_path)
    if not recs:
        print(f"ABORT: {in_path} has no usable records", file=sys.stderr)
        return 1

    shards_dir = Path(a.shards_dir)
    gold_cache: dict[int, str] = {}
    matrix: Counter = Counter()             # (code_label, opus_label)
    code_total: Counter = Counter()
    opus_total: Counter = Counter()
    by_arm: dict[str, Counter] = {}
    by_model: dict[str, Counter] = {}
    by_variant: dict[str, Counter] = {}
    residue: list[dict] = []

    t0 = time.time()
    n_skip_no_gold = 0
    n_skip_no_fixture_match = 0
    with open(out_jsonl, "w") as out_log:
        for r in recs:
            if r.get("arm") == "skip":
                continue
            i = r.get("i")
            fxr = by_i.get(i)
            if fxr is None:
                n_skip_no_fixture_match += 1
                continue
            if i not in gold_cache:
                gold_cache[i] = _gold(shards_dir, fxr.get("shard", ""),
                                      fxr["target_root"]) or ""
            gold = gold_cache[i]
            if not gold:
                n_skip_no_gold += 1
                continue
            q_asked = r.get("question_asked") or r.get("question_orig") or ""
            answer = r.get("answer") or ""
            v = code_judge(q_asked, answer, gold)
            opus_label = r.get("verdict", "?")
            matrix[(v.label, opus_label)] += 1
            code_total[v.label] += 1
            opus_total[opus_label] += 1
            by_arm.setdefault(r.get("arm", "?"), Counter())[v.label] += 1
            by_model.setdefault(r.get("model", "?"), Counter())[v.label] += 1
            by_variant.setdefault(r.get("variant", "?"), Counter())[v.label] += 1
            row = {**r,
                   "code_verdict": v.label,
                   "code_rationale": v.rationale,
                   "code_model": v.model,
                   "code_prompt_id": v.prompt_id,
                   "code_decision": v.decision}
            out_log.write(json.dumps(row, default=str) + "\n")
            out_log.flush()
            if v.label == "JUDGE_ERROR":
                residue.append({**r,
                                "code_rationale": v.rationale})

    dt = time.time() - t0
    n = sum(matrix.values())
    print(f"  re-graded {n} records in {dt:.1f}s "
          f"({n/dt:.0f} rec/s) — "
          f"skip(no_gold)={n_skip_no_gold} "
          f"skip(no_fixture_match)={n_skip_no_fixture_match}")
    print(f"  per-record JSONL: {out_jsonl}")

    L: list[str] = [
        f"# Code-judge re-scorecard for `{in_path.name}`",
        "",
        f"Generated {time.strftime('%Y-%m-%d %H:%M UTC', time.gmtime())}.",
        "",
        f"- Input sweep JSONL: `{in_path}` ({n} records re-graded; "
        f"skipped {n_skip_no_gold} no-gold + "
        f"{n_skip_no_fixture_match} fixture-miss).",
        f"- Fixture: `{a.fixture}` (must match the sweep's --fixture).",
        f"- Shards: `{shards_dir}`",
        f"- Code judge: `{CODE_JUDGE_MODEL}` "
        "(see `bench/judge_code.py`).",
        "",
        "**No LLM calls were made by this rescore.** The code-judge "
        "labels are deterministic; the original (Opus) labels are the "
        "ones already on disk in the sweep JSONL — they are not "
        "re-run here. JUDGE_ERROR residue is the eventual "
        "LLM-batch needle-haystack input (see fox 2026-05-19).",
        "",
    ]
    L += _agree_md(matrix, code_total, opus_total)
    L += [""]
    L += _scorecard_md(by_arm, "Code-judge scorecard by arm")
    L += [""]
    L += _scorecard_md(by_model, "Code-judge scorecard by model")
    L += [""]
    L += _scorecard_md(by_variant, "Code-judge scorecard by variant")
    L += [""]
    L += _residue_md(residue)

    out_md.parent.mkdir(parents=True, exist_ok=True)
    out_md.write_text("\n".join(L) + "\n")
    print(f"  markdown scorecard: {out_md}")
    return 0


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