"""Measure an LLM judge against human labels. Python 3.11+, standard library only.

Usage: python3 judge_agreement.py LABELS.csv   (or --demo)
LABELS.csv has the header case_id,human_label,judge_label; each label is pass or fail.
Rows with a human label must come from the test split, never used to write the prompt.
Rows with a blank human label were graded by the judge alone; their pass rate is corrected.
Exit status 1: the judge missed a bar, or the input can't be used.
"""
import csv
import math
import sys
from statistics import NormalDist

MIN_TPR, MIN_TNR = 0.80, 0.90  # agreed before the run; each interval's lower end must reach it
Z = NormalDist().inv_cdf(0.975)  # 95% intervals


def wilson(k, n):
    centre = (k + Z * Z / 2) / (n + Z * Z)
    half = Z / (n + Z * Z) * math.sqrt(k * (n - k) / n + Z * Z / 4)
    return max(0.0, centre - half), min(1.0, centre + half)


def clip(x):
    return min(1.0, max(0.0, x))


def corrected(passed, n, tp, n_fail, tn, n_pass):
    """Rogan-Gladen pass rate, with the adjusted 95% interval of Lee et al. (arXiv:2511.21140)."""
    tpr, tnr = tp / n_fail, tn / n_pass
    estimate = (passed / n + tpr - 1) / (tpr + tnr - 1)
    p = (passed + Z * Z / 2) / (n + Z * Z)  # their eq. 7: pseudo-counts on all three rates
    q0, q1 = (tp + 1) / (n_fail + 2), (tn + 1) / (n_pass + 2)
    if q0 + q1 <= 1:  # too few labels to bound it
        return estimate, 0.0, 1.0
    n, m0, m1 = n + Z * Z, n_fail + 2, n_pass + 2
    theta = (p + q0 - 1) / (q0 + q1 - 1)
    shift = 2 * Z * Z * (-(1 - theta) * q0 * (1 - q0) / m0 + theta * q1 * (1 - q1) / m1)  # eq. 8
    se = math.sqrt(p * (1 - p) / n + (1 - theta) ** 2 * q0 * (1 - q0) / m0
                   + theta ** 2 * q1 * (1 - q1) / m1) / (q0 + q1 - 1)  # eq. 6
    return estimate, clip(theta + shift - Z * se), clip(theta + shift + Z * se)


def show(label, text):
    print(label.ljust(26) + text)


def rate(k, n):
    low, high = wilson(k, n)
    return f"{k}/{n} = {k / n:.1%} ({100 * low:.1f}-{100 * high:.1f})"


def report(rows):
    pairs, judge_only, seen = [], [], set()
    for number, row in enumerate(rows, start=1):
        case, human, judge = ((row.get(key) or "").strip()
                              for key in ("case_id", "human_label", "judge_label"))
        human, judge = human.lower(), judge.lower()
        if (not case or case in seen or judge not in ("pass", "fail")
                or human not in ("pass", "fail", "")):
            sys.exit(f"row {number}: need a new case_id, a judge_label of pass or fail "
                     "and a human_label of pass, fail or blank")
        seen.add(case)
        if human:
            pairs.append((human, judge))
        else:
            judge_only.append(judge)
    tp, fn = pairs.count(("fail", "fail")), pairs.count(("fail", "pass"))
    fp, tn = pairs.count(("pass", "fail")), pairs.count(("pass", "pass"))
    n_fail, n_pass, n = tp + fn, fp + tn, len(pairs)
    if not (n_fail and n_pass):
        sys.exit("the labelled rows need at least one human fail and one human pass")
    print(f"labelled rows: {n} (human labels: {n_fail} fail, {n_pass} pass)")
    print(f"{'':12}{'judge fail':>12}{'judge pass':>12}")
    print(f"{'human fail':12}{tp:>12}{fn:>12}")
    print(f"{'human pass':12}{fp:>12}{tn:>12}")
    ok = True
    for label, k, m, bar in (("TPR, failures caught:", tp, n_fail, MIN_TPR),
                             ("TNR, good outputs passed:", tn, n_pass, MIN_TNR)):
        met = wilson(k, m)[0] >= bar
        ok = ok and met
        show(label, f"{rate(k, m)}  bar {bar:.0%} at the lower end: {'met' if met else 'missed'}")
    p_o = (tp + tn) / n
    p_e = (n_fail * (tp + fp) + n_pass * (fn + tn)) / n ** 2
    show("raw agreement:", f"{tp + tn}/{n} = {p_o:.1%}")
    show("Cohen's kappa:", f"{(p_o - p_e) / (1 - p_e):.3f}")
    if judge_only:
        if tp / n_fail + tn / n_pass <= 1:
            sys.exit("TPR + TNR is 1 or less: the judge is no better than chance")
        passed, m = judge_only.count("pass"), len(judge_only)
        estimate, low, high = corrected(passed, m, tp, n_fail, tn, n_pass)
        print(f"judge-only rows: {m}")
        show("judge's raw pass rate:", rate(passed, m))
        show("corrected pass rate:", f"{clip(estimate):.1%} ({100 * low:.1f}-{100 * high:.1f})")
        if not 0 <= estimate <= 1:
            print("warning: estimate clipped; these error rates don't fit the judge-only rows")
    print("decision: " + ("the judge meets both bars" if ok
                          else "the judge missed a bar; don't let it gate a release yet"))
    return ok


def demo():
    """Made-up rows for one failure class: 120 test-split labels, 400 judge-only verdicts."""
    counts = {("fail", "fail"): 30, ("fail", "pass"): 10, ("pass", "fail"): 2,
              ("pass", "pass"): 78, ("", "fail"): 80, ("", "pass"): 320}
    labels = [pair for pair, count in counts.items() for _ in range(count)]
    return [{"case_id": f"tone-{i:03d}", "human_label": human, "judge_label": judge}
            for i, (human, judge) in enumerate(labels, start=1)]


if __name__ == "__main__":
    if sys.argv[1:] == ["--demo"]:
        all_rows = demo()
    elif len(sys.argv) == 2:
        with open(sys.argv[1], newline="", encoding="utf-8") as f:
            all_rows = list(csv.DictReader(f))  # report() stops on a missing column
    else:
        sys.exit(__doc__)
    sys.exit(0 if report(all_rows) else 1)