#!/usr/bin/env python3
"""Deterministic protocol simulator for bias-audited pairwise judging.

This is not an LLM benchmark.  It exposes the mechanics of position/verbosity
bias, order swapping, rubric scoring, abstention, and panel aggregation with
synthetic answer features whose gold preference is known.
"""

from __future__ import annotations

import argparse
import hashlib
import random
from dataclasses import dataclass
from typing import Iterable, Sequence


SEED = 61


@dataclass(frozen=True)
class Answer:
    correctness: int
    instruction: int
    evidence: int
    clarity: int
    tokens: int


@dataclass(frozen=True)
class Pair:
    pair_id: str
    a: Answer
    b: Answer


@dataclass(frozen=True)
class Judge:
    name: str
    position_bias: float
    verbosity_bias: float
    noise: float


def rubric_score(answer: Answer) -> float:
    """Task-specific score; factual correctness is intentionally dominant."""
    return (
        0.45 * answer.correctness
        + 0.25 * answer.instruction
        + 0.20 * answer.evidence
        + 0.10 * answer.clarity
    )


def stable_noise(key: str, scale: float) -> float:
    digest = hashlib.sha256(key.encode("utf-8")).digest()
    unit = int.from_bytes(digest[:8], "big") / (2**64 - 1)
    return (2.0 * unit - 1.0) * scale


def make_pairs(n: int = 120) -> list[Pair]:
    rng = random.Random(SEED)
    pairs: list[Pair] = []
    while len(pairs) < n:
        answers = []
        for _ in range(2):
            answers.append(
                Answer(
                    correctness=rng.randint(0, 4),
                    instruction=rng.randint(0, 4),
                    evidence=rng.randint(0, 4),
                    clarity=rng.randint(0, 4),
                    tokens=rng.randint(40, 420),
                )
            )
        a, b = answers
        # Remove near-ties so the synthetic gold label is unambiguous.
        if abs(rubric_score(a) - rubric_score(b)) < 0.28:
            continue
        pairs.append(Pair(f"p{len(pairs):03d}", a, b))
    return pairs


def gold(pair: Pair) -> str:
    return "A" if rubric_score(pair.a) > rubric_score(pair.b) else "B"


def judge_once(pair: Pair, judge: Judge, order: str, rubric: bool) -> str:
    first, second = (pair.a, pair.b) if order == "AB" else (pair.b, pair.a)
    if rubric:
        first_score = rubric_score(first)
        second_score = rubric_score(second)
        length_weight = judge.verbosity_bias * 0.20
    else:
        # A vague "overall quality" score overweights polish and length.
        first_score = 0.55 * rubric_score(first) + 0.45 * first.clarity
        second_score = 0.55 * rubric_score(second) + 0.45 * second.clarity
        length_weight = judge.verbosity_bias

    first_score += length_weight * first.tokens / 420.0 + judge.position_bias
    second_score += length_weight * second.tokens / 420.0
    delta = first_score - second_score
    delta += stable_noise(f"{pair.pair_id}:{judge.name}:{order}:{rubric}", judge.noise)
    slot_winner = "FIRST" if delta > 0 else "SECOND"
    if order == "AB":
        return "A" if slot_winner == "FIRST" else "B"
    return "B" if slot_winner == "FIRST" else "A"


def accuracy(predictions: Sequence[str], labels: Sequence[str]) -> float:
    return sum(p == y for p, y in zip(predictions, labels)) / len(labels)


def evaluate_single(pairs: Sequence[Pair], judge: Judge, rubric: bool) -> dict[str, float]:
    labels = [gold(pair) for pair in pairs]
    ab = [judge_once(pair, judge, "AB", rubric) for pair in pairs]
    ba = [judge_once(pair, judge, "BA", rubric) for pair in pairs]
    consistent = [x == y for x, y in zip(ab, ba)]
    accepted = [x for x, ok in zip(ab, consistent) if ok]
    accepted_gold = [y for y, ok in zip(labels, consistent) if ok]
    return {
        "ab_accuracy": accuracy(ab, labels),
        "ba_accuracy": accuracy(ba, labels),
        "swap_consistency": sum(consistent) / len(consistent),
        "balanced_coverage": len(accepted) / len(pairs),
        "balanced_accuracy": accuracy(accepted, accepted_gold),
    }


def majority(votes: Iterable[str]) -> str:
    votes = list(votes)
    return "A" if votes.count("A") > votes.count("B") else "B"


def evaluate_panel(pairs: Sequence[Pair], judges: Sequence[Judge]) -> dict[str, float]:
    labels = [gold(pair) for pair in pairs]
    predictions: list[str] = []
    unanimous_order_stable = 0
    for pair in pairs:
        per_judge = []
        stable = True
        for judge in judges:
            ab = judge_once(pair, judge, "AB", rubric=True)
            ba = judge_once(pair, judge, "BA", rubric=True)
            per_judge.append(ab if ab == ba else majority([ab, ba]))
            stable &= ab == ba
        predictions.append(majority(per_judge))
        unanimous_order_stable += stable
    return {
        "accuracy": accuracy(predictions, labels),
        "all_judges_swap_stable": unanimous_order_stable / len(pairs),
    }


def run_checks() -> None:
    pairs = make_pairs()
    biased = Judge("biased", position_bias=0.70, verbosity_bias=1.10, noise=0.22)
    generic = evaluate_single(pairs, biased, rubric=False)
    rubric = evaluate_single(pairs, biased, rubric=True)
    panel = evaluate_panel(
        pairs,
        [
            Judge("cyan", 0.26, 0.30, 0.18),
            Judge("amber", -0.18, 0.12, 0.24),
            Judge("green", 0.08, -0.05, 0.20),
        ],
    )

    print(f"seed={SEED} pairs={len(pairs)}")
    print(
        "generic "
        f"AB_acc={generic['ab_accuracy']:.3f} BA_acc={generic['ba_accuracy']:.3f} "
        f"swap_consistency={generic['swap_consistency']:.3f} "
        f"balanced_coverage={generic['balanced_coverage']:.3f} "
        f"balanced_acc={generic['balanced_accuracy']:.3f}"
    )
    print(
        "rubric  "
        f"AB_acc={rubric['ab_accuracy']:.3f} BA_acc={rubric['ba_accuracy']:.3f} "
        f"swap_consistency={rubric['swap_consistency']:.3f} "
        f"balanced_coverage={rubric['balanced_coverage']:.3f} "
        f"balanced_acc={rubric['balanced_accuracy']:.3f}"
    )
    print(
        "panel   "
        f"accuracy={panel['accuracy']:.3f} "
        f"all_judges_swap_stable={panel['all_judges_swap_stable']:.3f}"
    )

    assert generic["swap_consistency"] < 0.80
    assert rubric["balanced_accuracy"] >= rubric["ab_accuracy"]
    assert rubric["balanced_accuracy"] >= rubric["ba_accuracy"]
    assert rubric["balanced_accuracy"] > generic["balanced_accuracy"]
    assert panel["accuracy"] > generic["ab_accuracy"]
    print("checks=passed")


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--check-only", action="store_true", help="run deterministic assertions")
    args = parser.parse_args()
    if not args.check_only:
        parser.error("this teaching script only supports --check-only")
    run_checks()


if __name__ == "__main__":
    main()