"""Week 8a lab: graders, a mock LLM judge with injectable biases, bias measurement and fixes, pointwise and
pairwise scoring with Bradley-Terry, judge-human agreement and calibration, pass@k and pass^k with a noisy
grader, an AutoMetrics-style metric induction toy, and a small eval harness with caching, retries, seeds and
a result log. Standard library only. Deterministic. No network, no model.

Everything here is made up: the questions, the "responses", the human ratings, the judges. The judges are
not language models. Each is a rule with noise and a few bias knobs. Where a method follows a paper, the
text near it says which and says what I changed.
"""
import hashlib
import json
import math
import random
import re
from collections import Counter, defaultdict

# ---------------------------------------------------------------- part 0: shared helpers


def mean(xs):
    xs = list(xs)
    return sum(xs) / len(xs) if xs else 0.0


def stdev(xs, ddof=1):
    xs = list(xs)
    n = len(xs)
    if n - ddof <= 0:
        return 0.0
    m = mean(xs)
    return math.sqrt(sum((x - m) ** 2 for x in xs) / (n - ddof))


def sigmoid(x):
    if x >= 0:
        return 1.0 / (1.0 + math.exp(-x))
    e = math.exp(x)
    return e / (1.0 + e)


def clip(x, lo, hi):
    return max(lo, min(hi, x))


def normalize(text):
    """Lower-case, drop punctuation except digits, dots and minus, collapse spaces."""
    t = re.sub(r"[^a-z0-9.\- ]+", " ", str(text).lower())
    return re.sub(r"\s+", " ", t).strip().rstrip(".").strip()


def ranks(xs):
    """Average ranks (1 is the smallest). Ties share the mean of their ranks."""
    order = sorted(range(len(xs)), key=lambda i: xs[i])
    out = [0.0] * len(xs)
    i = 0
    while i < len(order):
        j = i
        while j + 1 < len(order) and xs[order[j + 1]] == xs[order[i]]:
            j += 1
        r = (i + j) / 2.0 + 1.0
        for k in range(i, j + 1):
            out[order[k]] = r
        i = j + 1
    return out


def pearson(x, y):
    mx, my = mean(x), mean(y)
    sx = math.sqrt(sum((a - mx) ** 2 for a in x))
    sy = math.sqrt(sum((b - my) ** 2 for b in y))
    if sx == 0 or sy == 0:
        return 0.0
    return sum((a - mx) * (b - my) for a, b in zip(x, y)) / (sx * sy)


def spearman(x, y):
    """Pearson correlation of the average ranks."""
    return pearson(ranks(x), ranks(y))


def kendall_tau_b(x, y):
    """Kendall's tau-b, which corrects for ties in either list."""
    n = len(x)
    conc = disc = tx = ty = 0
    for i in range(n):
        for j in range(i + 1, n):
            dx, dy = x[i] - x[j], y[i] - y[j]
            if dx == 0 and dy == 0:
                continue
            if dx == 0:
                tx += 1
            elif dy == 0:
                ty += 1
            elif dx * dy > 0:
                conc += 1
            else:
                disc += 1
    den = math.sqrt((conc + disc + tx) * (conc + disc + ty))
    return (conc - disc) / den if den else 0.0


def comb(n, k):
    if k < 0 or k > n:
        return 0
    return math.comb(n, k)


def pass_at_k(n, c, k):
    """Unbiased pass@k from c passes in n trials (Chen et al.): 1 - C(n-c, k) / C(n, k)."""
    if k > n:
        raise ValueError("k > n")
    return 1.0 - comb(n - c, k) / comb(n, k)


def pass_hat_k(n, c, k):
    """Unbiased pass^k from c passes in n trials (tau-bench): C(c, k) / C(n, k)."""
    if k > n:
        raise ValueError("k > n")
    return comb(c, k) / comb(n, k)


def wilson(k, n, z=1.96):
    if n == 0:
        return (0.0, 1.0)
    p = k / n
    den = 1 + z * z / n
    c = (p + z * z / (2 * n)) / den
    h = z * math.sqrt(p * (1 - p) / n + z * z / (4 * n * n)) / den
    return (max(0.0, c - h), min(1.0, c + h))


# ---------------------------------------------------------------- part 1: code-based graders
# The first of the three grader types. Cheap, fast, repeatable, and blind to anything that was not written down.


def exact_match(answer, gold):
    return normalize(answer) == normalize(gold)


def contains_all(answer, needles):
    a = normalize(answer)
    return all(normalize(n) in a for n in needles)


def regex_match(answer, pattern):
    return re.search(pattern, str(answer)) is not None


def first_number(text):
    m = re.search(r"-?\d+(?:\.\d+)?", str(text))
    return float(m.group(0)) if m else None


def number_close(answer, gold, tol=1e-6):
    """The first number in the answer is within tol of the gold number."""
    v = first_number(answer)
    return v is not None and abs(v - gold) <= tol


def json_fields(answer, required):
    """The answer is a JSON object and every key in `required` has that value."""
    try:
        obj = json.loads(answer)
    except (ValueError, TypeError):
        return False
    return isinstance(obj, dict) and all(k in obj and obj[k] == v for k, v in required.items())


def tool_calls_ok(calls, required, forbidden=()):
    """`required` appears in `calls` as a subsequence, and no forbidden tool name is used."""
    names = [c[0] if isinstance(c, (tuple, list)) else c for c in calls]
    if any(n in forbidden for n in names):
        return False
    it = iter(names)
    return all(any(n == r for n in it) for r in required)


def run_assertions(answer, assertions):
    """assertions: list of (kind, arg). Returns (share passed, all passed, list of results)."""
    fns = {"exact": exact_match, "contains": contains_all, "regex": regex_match, "number": number_close}
    res = []
    for kind, arg in assertions:
        if kind not in fns:
            raise ValueError("unknown assertion " + str(kind))
        res.append(bool(fns[kind](answer, arg)))
    return (mean(res) if res else 0.0, all(res), res)


# ---------------------------------------------------------------- part 2: a made-up world of responses
# Five "systems" answer 60 questions. Each response has hidden traits. A hidden quality number h comes from the
# traits. Humans see h with a little noise. The mock judges see h with more noise and with bias.

SYSTEMS = {
    # p_correct, p_steps, p_source, p_hedge, padding sentences on average
    "alpha": (0.85, 0.60, 0.40, 0.10, 0.0),
    "beta": (0.75, 0.50, 0.30, 0.15, 1.0),
    "gamma": (0.65, 0.50, 0.20, 0.20, 2.0),
    "delta": (0.55, 0.30, 0.20, 0.30, 0.0),
    "epsilon": (0.45, 0.30, 0.10, 0.40, 4.0),
}
FILLER = [
    "It is worth noting that this matters in practice.",
    "There are many things to consider here.",
    "This is a common question that comes up often.",
    "In general, the details can vary from case to case.",
    "As a reminder, it helps to think about this carefully.",
]
CUTS = (0.2, 0.6, 1.0, 1.4)  # latent quality to a 1..5 rating
PASS_THRESHOLD = 0.9  # a response passes for the human when h is at least this


def rating_of(h):
    return 1 + sum(1 for c in CUTS if h >= c)


def make_world(seed=0, n_q=60):
    rng = random.Random(f"world-{seed}")
    resps = []
    questions = []
    for q in range(n_q):
        gold = rng.randint(10, 99)
        diff = rng.gauss(0, 0.5)
        questions.append({"qid": q, "gold": gold, "difficulty": diff})
        for s, (pc, ps, pso, ph, pad) in SYSTEMS.items():
            correct = rng.random() < sigmoid(math.log(pc / (1 - pc)) - diff)
            steps = rng.random() < ps
            source = rng.random() < pso
            hedge = rng.random() < ph
            polite = rng.random() < 0.5
            n_pad = max(0, int(round(rng.gauss(pad, 0.7)))) if pad > 0 else 0
            bits = []
            if hedge:
                bits.append("I think maybe")
            if steps:
                bits.append("First, read the question. Then, work it out step by step because that avoids slips.")
            ans = gold if correct else gold + rng.choice([-7, -3, 3, 7, 11])
            bits.append(f"The answer is {ans}.")
            if source:
                bits.append("According to the docs, this matches the reference.")
            for _ in range(n_pad):
                bits.append(rng.choice(FILLER))
            if polite:
                bits.append("Happy to help!")
            text = " ".join(bits)
            h = (1.0 * correct + 0.4 * steps + 0.3 * source - 0.3 * hedge) + rng.gauss(0, 0.15)
            resps.append({"id": f"q{q}-{s}", "qid": q, "system": s, "text": text, "words": len(text.split()),
                          "traits": {"correct": correct, "steps": steps, "source": source, "hedge": hedge,
                                     "polite": polite, "pad": n_pad}, "h": h})
    ws = [r["words"] for r in resps]
    m, sd = mean(ws), stdev(ws)
    for r in resps:
        r["zlen"] = (r["words"] - m) / sd
    return {"seed": seed, "questions": questions, "responses": resps, "len_mean": m, "len_sd": sd,
            "by_q": _group(resps, "qid"), "systems": list(SYSTEMS)}


def _group(resps, key):
    out = defaultdict(dict)
    for r in resps:
        out[r[key]][r["system"]] = r
    return dict(out)


def pad_response(world, r, n_sentences=4):
    """The same response with filler added. Same quality, longer. Used to test for a taste for length."""
    extra = " ".join(FILLER[i % len(FILLER)] for i in range(n_sentences))
    words = r["words"] + len(extra.split())
    q = dict(r)
    q.update({"id": r["id"] + "+pad", "text": r["text"] + " " + extra, "words": words,
              "zlen": (words - world["len_mean"]) / world["len_sd"]})
    return q


def human_rating(r):
    return rating_of(r["h"])


def human_pass(r, thr=PASS_THRESHOLD):
    return r["h"] >= thr


def human_prefers(a, b, noise=0.12, seed=0):
    """A human's pairwise choice: 'A', 'B' or 'tie'. A little noise, and a tie when they are very close."""
    rng = random.Random(f"human-{seed}-{a['id']}-{b['id']}")
    d = (a["h"] - b["h"]) + rng.gauss(0, noise)
    if abs(d) < 0.05:
        return "tie"
    return "A" if d > 0 else "B"


def true_system_quality(world):
    by = defaultdict(list)
    for r in world["responses"]:
        by[r["system"]].append(r["h"])
    return {s: mean(v) for s, v in by.items()}


# ---------------------------------------------------------------- part 3: a mock judge with injectable biases


class MockJudge:
    """What the judge "sees" of a response is h plus a persistent error for this judge and response, plus a fresh
    error for each call, plus bias terms. It never reads the text. Bias knobs:
      pos_bias  added to the response shown first in a pairwise call (position bias)
      len_bias  added per standard deviation of length (verbosity bias)
      self_bias added when the response comes from `self_system` (self-preference)
    `k` is the sample index: the same k gives the same answer, a different k draws a new call error."""

    def __init__(self, name="judge", sigma=0.25, item_noise=0.15, pos_bias=0.0, len_bias=0.0, self_bias=0.0,
                 self_system=None, tie_margin=0.1, seed=0, taste=None):
        self.taste = taste
        self.name, self.sigma, self.item_noise = name, sigma, item_noise
        self.pos_bias, self.len_bias, self.self_bias = pos_bias, len_bias, self_bias
        self.self_system, self.tie_margin, self.seed = self_system, tie_margin, seed

    def _gauss(self, tag, sd):
        if sd <= 0:
            return 0.0
        return random.Random(f"{self.name}|{self.seed}|{tag}").gauss(0, sd)

    def _taste_value(self, r):
        """A judge that does not see h. It scores the traits with its own weights, which differ from the human's."""
        t = r["traits"]
        return sum(w * (r["zlen"] if k == "zlen" else float(t[k])) for k, w in self.taste.items())

    def perceived(self, r, k=0, call=""):
        base = r["h"] if self.taste is None else self._taste_value(r)
        v = base + self._gauss("item|" + r["id"], self.item_noise) + self._gauss(f"call|{r['id']}|{call}|{k}", self.sigma)
        v += self.len_bias * r["zlen"]
        if self.self_system is not None and r["system"] == self.self_system:
            v += self.self_bias
        return v

    def pointwise(self, r, k=0):
        """A 1 to 5 score against a fixed rubric (the cut points are the human's)."""
        return rating_of(self.perceived(r, k, "pw"))

    def pairwise(self, a, b, k=0):
        """'A', 'B' or 'tie'. a is shown first."""
        call = f"pair|{a['id']}|{b['id']}"
        d = self.perceived(a, k, call) - self.perceived(b, k, call) + self.pos_bias
        if abs(d) < self.tie_margin:
            return "tie"
        return "A" if d > 0 else "B"

    def pass_confidence(self, r, thr=PASS_THRESHOLD, scale=0.35, k=0):
        """A made-up "probability that this passes": a sigmoid of the distance from the threshold."""
        return sigmoid((self.perceived(r, k, "pc") - thr) / scale)

    def pass_label(self, r, thr=PASS_THRESHOLD, k=0):
        return self.perceived(r, k, "pl") >= thr


# ---------------------------------------------------------------- part 4: measuring the biases


def pair_list(world, n_pairs=None, seed=0):
    """All pairs of different systems on the same question, as (a, b) response pairs."""
    pairs = []
    for q, d in sorted(world["by_q"].items()):
        names = sorted(d)
        for i in range(len(names)):
            for j in range(i + 1, len(names)):
                pairs.append((d[names[i]], d[names[j]]))
    if n_pairs is not None:
        rng = random.Random(f"pairs-{seed}")
        pairs = rng.sample(pairs, min(n_pairs, len(pairs)))
    return pairs


def position_report(judge, pairs, k=0):
    """Swap the order of each pair. Consistent means the same response wins (or a tie) both times.
    first_rate is the share of non-tie verdicts, over both orders, that pick the response shown first."""
    cons = first = nontie = 0
    for a, b in pairs:
        v1 = judge.pairwise(a, b, k)
        v2 = judge.pairwise(b, a, k)
        w1 = {"A": a["id"], "B": b["id"], "tie": None}[v1]
        w2 = {"A": b["id"], "B": a["id"], "tie": None}[v2]
        cons += int(w1 == w2)
        for v in (v1, v2):
            if v != "tie":
                nontie += 1
                first += int(v == "A")
    n = len(pairs)
    return {"consistency": cons / n if n else 0.0, "first_rate": first / nontie if nontie else 0.0, "pairs": n}


def verbosity_report(judge, world, responses, n_sentences=4, k=0):
    """Compare each response with a padded copy of itself, in both orders. Equal quality, so an unbiased judge
    gives the padded copy half of the credit (a tie counts as half). Returns the padded copy's share."""
    credit = 0.0
    n = 0
    for r in responses:
        p = pad_response(world, r, n_sentences)
        for a, b in ((p, r), (r, p)):
            v = judge.pairwise(a, b, k)
            if v == "tie":
                credit += 0.5
            elif (v == "A") == (a is p):
                credit += 1.0
            n += 1
    return {"padded_share": credit / n if n else 0.0, "n": n}


def self_preference_report(judge, world, pairs=None, k=0):
    """For pairs that include the judge's own system: the judge's win rate for the own side (both orders, a tie is
    half) minus the human win rate (from h, no noise, a tie is half). Positive means the judge likes its own."""
    s = judge.self_system
    pairs = pairs if pairs is not None else pair_list(world)
    jw = hw = n = 0.0
    for a, b in pairs:
        if s not in (a["system"], b["system"]):
            continue
        own, other = (a, b) if a["system"] == s else (b, a)
        for x, y in ((own, other), (other, own)):
            v = judge.pairwise(x, y, k)
            if v == "tie":
                jw += 0.5
            elif (v == "A") == (x is own):
                jw += 1.0
            n += 1
        hw += 1.0 if own["h"] > other["h"] else 0.0
    hn = n / 2.0
    return {"judge_winrate": jw / n if n else 0.0, "human_winrate": hw / hn if hn else 0.0,
            "gap": (jw / n - hw / hn) if n else 0.0, "n": int(n)}


def length_correlation(judge, responses, k=0):
    """Correlation between the judge's pointwise score and the number of words, minus the same for humans."""
    j = [judge.pointwise(r, k) for r in responses]
    h = [human_rating(r) for r in responses]
    w = [r["words"] for r in responses]
    return {"judge": pearson(j, w), "human": pearson(h, w)}


# ---------------------------------------------------------------- part 5: mitigations


def swapped_verdict(judge, a, b, k=0):
    """Ask in both orders. Declare a winner only if both orders agree. Else a tie. (Zheng et al. call this the
    conservative approach to position bias.) Returns 'a', 'b' or 'tie' for the responses, not for the slots."""
    v1 = judge.pairwise(a, b, k)
    v2 = judge.pairwise(b, a, k)
    w1 = "a" if v1 == "A" else "b" if v1 == "B" else "tie"
    w2 = "b" if v2 == "A" else "a" if v2 == "B" else "tie"
    return w1 if w1 == w2 else "tie"


def random_order_verdict(judge, a, b, rng, k=0):
    """Randomize which response is shown first. Unbiased on average, noisier for any one pair."""
    if rng.random() < 0.5:
        v = judge.pairwise(a, b, k)
        return {"A": "a", "B": "b", "tie": "tie"}[v]
    v = judge.pairwise(b, a, k)
    return {"A": "b", "B": "a", "tie": "tie"}[v]


def panel_verdict(judges, a, b, k=0, swap=True, exclude_own=True):
    """A panel of judges. Each votes (in both orders if swap). The sign of the vote total decides. With exclude_own,
    a judge does not vote on a pair that includes its own system."""
    total = 0
    voters = 0
    for j in judges:
        if exclude_own and j.self_system in (a["system"], b["system"]):
            continue
        voters += 1
        v = swapped_verdict(j, a, b, k) if swap else {"A": "a", "B": "b", "tie": "tie"}[j.pairwise(a, b, k)]
        total += {"a": 1, "b": -1, "tie": 0}[v]
    if voters == 0 or total == 0:
        return "tie"
    return "a" if total > 0 else "b"


def fit_logistic(X, y, l2=1e-3, iters=200, lr=0.5):
    """Logistic regression by gradient descent, with an intercept. X is a list of feature lists, y in [0, 1]
    (half credit allowed). Returns [intercept, coef...]."""
    n, d = len(X), len(X[0]) if X else 0
    w = [0.0] * (d + 1)
    for _ in range(iters):
        g = [0.0] * (d + 1)
        for xi, yi in zip(X, y):
            p = sigmoid(w[0] + sum(w[j + 1] * xi[j] for j in range(d)))
            e = p - yi
            g[0] += e
            for j in range(d):
                g[j + 1] += e * xi[j]
        for j in range(d + 1):
            reg = l2 * w[j] if j else 0.0
            w[j] -= lr * (g[j] / n + reg)
    return w


def length_controlled_winrate(rows, **kw):
    """rows: dicts with 'len_diff' (standardized length of the first minus the second) and 'pref' (1, 0.5 or 0 for
    the first). Fit pref ~ sigmoid(b0 + b1 * len_diff) and report the prediction at len_diff = 0. This is a
    simplified version of the idea in Length-Controlled AlpacaEval, which uses a richer model."""
    X = [[r["len_diff"]] for r in rows]
    y = [r["pref"] for r in rows]
    w = fit_logistic(X, y, **kw)
    raw = mean(y)
    return {"raw": raw, "controlled": sigmoid(w[0]), "length_coef": w[1]}


def system_pair_rows(world, judge, system, baseline, k=0, both_orders=True):
    """Pairwise rows for `system` against `baseline` on every question, with the length difference."""
    rows = []
    for q, d in sorted(world["by_q"].items()):
        a, b = d[system], d[baseline]
        orders = ((a, b), (b, a)) if both_orders else ((a, b),)
        for x, y in orders:
            v = judge.pairwise(x, y, k)
            first = 1.0 if v == "A" else 0.0 if v == "B" else 0.5
            mine = first if x is a else 1.0 - first
            rows.append({"len_diff": a["zlen"] - b["zlen"], "pref": mine})
    return rows


# ---------------------------------------------------------------- part 6: pointwise and pairwise rankings


def comparisons_from(world, verdict_fn):
    """verdict_fn(a, b) returns 'a', 'b' or 'tie'. Returns (system_a, system_b, outcome for a: 1, 0.5, 0)."""
    out = []
    for a, b in pair_list(world):
        v = verdict_fn(a, b)
        out.append((a["system"], b["system"], {"a": 1.0, "tie": 0.5, "b": 0.0}[v]))
    return out


def bradley_terry(comps, systems, iters=500, tol=1e-9, prior=0.5):
    """Bradley-Terry strengths by the standard minorization-maximization update. A tie is half a win for each side.
    `prior` adds that many wins of each side against each other side, so that a system with no losses stays finite.
    Returns log-strengths with mean 0."""
    w = {s: 0.0 for s in systems}
    n = defaultdict(float)
    for a, b, o in comps:
        w[a] += o
        w[b] += 1 - o
        n[(a, b)] += 1
        n[(b, a)] += 1
    for s in systems:
        for t in systems:
            if s != t:
                w[s] += prior
                n[(s, t)] += 2 * prior
    p = {s: 1.0 for s in systems}
    for _ in range(iters):
        new = {}
        for s in systems:
            den = sum(n[(s, t)] / (p[s] + p[t]) for t in systems if t != s and n[(s, t)] > 0)
            new[s] = w[s] / den if den else p[s]
        g = math.exp(mean(math.log(v) for v in new.values()))
        new = {s: v / g for s, v in new.items()}
        delta = max(abs(math.log(new[s]) - math.log(p[s])) for s in systems)
        p = new
        if delta < tol:
            break
    return {s: math.log(v) for s, v in p.items()}


def to_elo(log_strengths, base=1000.0):
    """Put log-strengths on the familiar scale: 400 points per factor of 10 in odds."""
    c = 400.0 / math.log(10)
    return {s: base + c * v for s, v in log_strengths.items()}


def elo_online(comps, k=32.0, seed=0, base=1000.0):
    """Online Elo over the comparisons in a random order. The answer depends on the order, which is one reason
    to fit all the comparisons at once, as bradley_terry does."""
    rng = random.Random(f"elo-{seed}")
    order = list(comps)
    rng.shuffle(order)
    r = defaultdict(lambda: base)
    for a, b, o in order:
        ea = 1.0 / (1.0 + 10 ** ((r[b] - r[a]) / 400.0))
        r[a] += k * (o - ea)
        r[b] += k * ((1 - o) - (1 - ea))
    return dict(r)


def bootstrap_bt(comps, systems, n_boot=100, seed=0):
    """Resample the comparisons with replacement and refit. Returns per-system (low, high) of the 95% interval on
    the log-strength, and how often each system lands at each rank."""
    rng = random.Random(f"boot-{seed}")
    fits = []
    for _ in range(n_boot):
        samp = [comps[rng.randrange(len(comps))] for _ in comps]
        fits.append(bradley_terry(samp, systems, iters=100))
    ci = {}
    rank_counts = {s: Counter() for s in systems}
    for f in fits:
        order = sorted(systems, key=lambda s: -f[s])
        for pos, s in enumerate(order):
            rank_counts[s][pos + 1] += 1
    for s in systems:
        v = sorted(f[s] for f in fits)
        ci[s] = (v[int(0.025 * (n_boot - 1))], v[int(math.ceil(0.975 * (n_boot - 1)))])
    return {"ci": ci, "rank_counts": rank_counts}


def pointwise_system_scores(world, judge, k=0):
    by = defaultdict(list)
    for r in world["responses"]:
        by[r["system"]].append(judge.pointwise(r, k))
    return {s: mean(v) for s, v in by.items()}


def ranking_agreement(scores, truth):
    keys = sorted(truth)
    a = [scores[s] for s in keys]
    b = [truth[s] for s in keys]
    return {"kendall": kendall_tau_b(a, b), "spearman": spearman(a, b)}


# ---------------------------------------------------------------- part 7: judge and human agreement, calibration


def percent_agreement(a, b):
    return mean(1.0 if x == y else 0.0 for x, y in zip(a, b))


def cohen_kappa(a, b):
    n = len(a)
    po = percent_agreement(a, b)
    ca, cb = Counter(a), Counter(b)
    pe = sum((ca[c] / n) * (cb[c] / n) for c in set(ca) | set(cb))
    return (po - pe) / (1 - pe) if pe < 1 else 0.0


def weighted_kappa(a, b, levels=(1, 2, 3, 4, 5)):
    """Quadratic-weighted kappa for ordinal ratings."""
    n = len(a)
    k = len(levels)
    idx = {v: i for i, v in enumerate(levels)}
    obs = [[0.0] * k for _ in range(k)]
    for x, y in zip(a, b):
        obs[idx[x]][idx[y]] += 1.0 / n
    ra = [sum(obs[i]) for i in range(k)]
    rb = [sum(obs[i][j] for i in range(k)) for j in range(k)]
    num = den = 0.0
    for i in range(k):
        for j in range(k):
            w = ((i - j) / (k - 1)) ** 2
            num += w * obs[i][j]
            den += w * ra[i] * rb[j]
    return 1 - num / den if den else 0.0


def confusion(judge_labels, human_labels, positive=False):
    """Treat `positive` as the class you want to catch (default: a fail, which is False). Returns tp, fp, tn, fn."""
    tp = fp = tn = fn = 0
    for j, h in zip(judge_labels, human_labels):
        jp, hp = (j == positive), (h == positive)
        if jp and hp:
            tp += 1
        elif jp and not hp:
            fp += 1
        elif not jp and not hp:
            tn += 1
        else:
            fn += 1
    return tp, fp, tn, fn


def tpr_tnr(judge_labels, human_labels, positive=False):
    tp, fp, tn, fn = confusion(judge_labels, human_labels, positive)
    return {"tpr": tp / (tp + fn) if tp + fn else 0.0, "tnr": tn / (tn + fp) if tn + fp else 0.0}


def reliability_bins(confs, labels, n_bins=10):
    """For each confidence bin: count, mean confidence, share of positives."""
    bins = [[] for _ in range(n_bins)]
    for c, y in zip(confs, labels):
        bins[min(n_bins - 1, int(c * n_bins))].append((c, 1.0 if y else 0.0))
    return [{"n": len(b), "conf": mean(c for c, _ in b), "freq": mean(y for _, y in b)} if b else
            {"n": 0, "conf": 0.0, "freq": 0.0} for b in bins]


def expected_calibration_error(confs, labels, n_bins=10):
    n = len(confs)
    return sum(b["n"] / n * abs(b["conf"] - b["freq"]) for b in reliability_bins(confs, labels, n_bins) if b["n"])


def platt_fit(scores, labels, **kw):
    """Fit sigmoid(a + b * score) to the human labels. Returns (a, b)."""
    w = fit_logistic([[s] for s in scores], [1.0 if y else 0.0 for y in labels], **kw)
    return w[0], w[1]


def score_to_pass_rate(scores, labels):
    """For each judge score, the share of responses that the human passes."""
    d = defaultdict(list)
    for s, y in zip(scores, labels):
        d[s].append(1.0 if y else 0.0)
    return {s: (mean(v), len(v)) for s, v in sorted(d.items())}


def imbalanced_labels(world, fail_rate):
    """Human pass labels, but with the threshold moved so that about `fail_rate` of responses fail."""
    hs = sorted(r["h"] for r in world["responses"])
    thr = hs[int(fail_rate * (len(hs) - 1))]
    return thr


# ---------------------------------------------------------------- part 8: pass@k and pass^k with a noisy grader


def graded_runs(task_p, trials, tpr=1.0, tnr=1.0, seed=0):
    """Each task has a true success rate. Run `trials` per task. A grader has a hit rate `tpr` (it passes a true
    success) and a correct-reject rate `tnr` (it fails a true failure). Returns (true matrix, judged matrix)."""
    truth, judged = [], []
    for i, p in enumerate(task_p):
        rng = random.Random(f"gr-{seed}-{i}")
        t_row, j_row = [], []
        for _ in range(trials):
            ok = rng.random() < p
            seen = (rng.random() < tpr) if ok else (rng.random() >= tnr)
            t_row.append(ok)
            j_row.append(seen)
        truth.append(t_row)
        judged.append(j_row)
    return truth, judged


def curves_from(mat, ks):
    n = len(mat[0])
    out = {"pass_at_k": [], "pass_hat_k": []}
    for k in ks:
        out["pass_at_k"].append(mean(pass_at_k(n, sum(row), k) for row in mat))
        out["pass_hat_k"].append(mean(pass_hat_k(n, sum(row), k) for row in mat))
    return out


def judge_rerun_agreement(judge, responses, runs=5, thr=PASS_THRESHOLD):
    """Ask the same judge the same question `runs` times (different call errors). Returns the share of responses on
    which every run gave the same pass/fail label, and the mean label flip rate against the first run."""
    same = 0
    flips = []
    for r in responses:
        labs = [judge.pass_label(r, thr, k) for k in range(runs)]
        same += int(len(set(labs)) == 1)
        flips.append(mean(1.0 if l != labs[0] else 0.0 for l in labs[1:]))
    return {"all_same": same / len(responses), "flip_rate": mean(flips)}


# ---------------------------------------------------------------- part 9: an AutoMetrics-style metric induction toy
# Generate candidate criteria, score every response on each (a noisy "LLM judge for one criterion"), z-normalize,
# regress onto a few human labels, keep the best few. This copies the shape of the AutoMetrics pipeline
# (generate, retrieve, regress, report). It does not call a model, it has no MetricBank, and it uses ridge
# regression, not the partial least squares of the paper.

CRITERIA = {
    # name: (what it looks at in the text, noise sd of the mock criterion judge)
    "shows_steps": (lambda r: 1.0 if "step by step" in r["text"] else 0.0, 0.20),
    "cites_a_source": (lambda r: 1.0 if "According to the docs" in r["text"] else 0.0, 0.20),
    "hedges": (lambda r: 1.0 if "I think maybe" in r["text"] else 0.0, 0.20),
    "polite_closing": (lambda r: 1.0 if "Happy to help" in r["text"] else 0.0, 0.20),
    "is_long": (lambda r: clip(r["words"] / 80.0, 0, 1), 0.10),
    "has_filler": (lambda r: 1.0 if any(f in r["text"] for f in FILLER) else 0.0, 0.20),
    "answer_looks_right": (lambda r: 1.0 if r["traits"]["correct"] else 0.0, 0.45),
    "mentions_answer": (lambda r: 1.0 if "The answer is" in r["text"] else 0.0, 0.05),
    "uses_because": (lambda r: 1.0 if "because" in r["text"] else 0.0, 0.20),
    "exclamation": (lambda r: 1.0 if "!" in r["text"] else 0.0, 0.10),
    "sentence_count": (lambda r: clip(r["text"].count(".") / 8.0, 0, 1), 0.10),
    "starts_with_then": (lambda r: 1.0 if "Then," in r["text"] else 0.0, 0.20),
}


TASTE = {"correct": 0.6, "steps": 0.1, "source": 0.0, "hedge": 0.0, "polite": 0.4, "zlen": 0.3}  # a judge whose taste differs from the human's


def criterion_scores(responses, names=None, seed=0):
    """A matrix (responses x criteria) of mock criterion-judge scores."""
    names = names or list(CRITERIA)
    X = []
    for r in responses:
        row = []
        for nm in names:
            fn, sd = CRITERIA[nm]
            noise = random.Random(f"crit|{seed}|{nm}|{r['id']}").gauss(0, sd) if sd else 0.0
            row.append(fn(r) + noise)
        X.append(row)
    return X


def zscore_cols(X, means=None, sds=None):
    d = len(X[0])
    means = means or [mean(row[j] for row in X) for j in range(d)]
    sds = sds or [stdev([row[j] for row in X]) or 1.0 for j in range(d)]
    return [[(row[j] - means[j]) / sds[j] for j in range(d)] for row in X], means, sds


def solve_linear(A, b):
    """Gaussian elimination with partial pivoting."""
    n = len(A)
    M = [list(A[i]) + [b[i]] for i in range(n)]
    for c in range(n):
        p = max(range(c, n), key=lambda i: abs(M[i][c]))
        if abs(M[p][c]) < 1e-12:
            raise ValueError("singular")
        M[c], M[p] = M[p], M[c]
        for i in range(c + 1, n):
            f = M[i][c] / M[c][c]
            for j in range(c, n + 1):
                M[i][j] -= f * M[c][j]
    x = [0.0] * n
    for i in range(n - 1, -1, -1):
        x[i] = (M[i][n] - sum(M[i][j] * x[j] for j in range(i + 1, n))) / M[i][i]
    return x


def ridge_fit(Z, y, lam=1.0):
    """Ridge regression on standardized columns with a centered target. Returns (coefs, y mean)."""
    d = len(Z[0])
    my = mean(y)
    A = [[sum(row[i] * row[j] for row in Z) + (lam if i == j else 0.0) for j in range(d)] for i in range(d)]
    b = [sum(row[i] * (yy - my) for row, yy in zip(Z, y)) for i in range(d)]
    return solve_linear(A, b), my


def induce_metric(responses, y, names=None, n_keep=5, lam=5.0, drop_negative=False, seed=0):
    """Fit on all candidate criteria, keep the n_keep with the largest |coefficient|, refit on those. With
    drop_negative, drop kept criteria whose refit coefficient is negative and refit once more.
    Returns a dict with names, weights, and a predict(responses) function."""
    names = names or list(CRITERIA)
    X = criterion_scores(responses, names, seed)
    Z, mu, sd = zscore_cols(X)
    coef, _ = ridge_fit(Z, y, lam)
    order = sorted(range(len(names)), key=lambda j: -abs(coef[j]))[:n_keep]
    keep = [names[j] for j in order]
    cols = order

    def refit(cols):
        Zs = [[row[j] for j in cols] for row in Z]
        return ridge_fit(Zs, y, lam)

    c2, my = refit(cols)
    if drop_negative and any(c < 0 for c in c2):
        cols = [j for j, c in zip(cols, c2) if c >= 0] or cols
        c2, my = refit(cols)
    keep = [names[j] for j in cols]

    def predict(rs):
        Xn = criterion_scores(rs, names, seed)
        out = []
        for row in Xn:
            out.append(my + sum(c * (row[j] - mu[j]) / sd[j] for c, j in zip(c2, cols)))
        return out

    return {"names": keep, "weights": c2, "predict": predict}


def induction_curve(world, sizes=(10, 20, 40, 80, 160), repeats=20, n_keep=5, seed=0):
    """Mean Kendall tau-b on held-out responses between the induced metric and the human rating, against the number
    of human labels. Also a baseline: one holistic mock judge's pointwise score on the same held-out set."""
    rs = world["responses"]
    base = MockJudge("holistic", sigma=0.25, item_noise=0.30, seed=seed)
    taste = MockJudge("holistic-taste", sigma=0.25, item_noise=0.30, seed=seed, taste=TASTE)
    rows = []
    for n in sizes:
        taus, btaus, wtaus = [], [], []
        for t in range(repeats):
            rng = random.Random(f"ind-{seed}-{n}-{t}")
            idx = list(range(len(rs)))
            rng.shuffle(idx)
            train = [rs[i] for i in idx[:n]]
            test = [rs[i] for i in idx[n:n + 120]]
            m = induce_metric(train, [human_rating(r) for r in train], n_keep=n_keep, seed=seed)
            taus.append(kendall_tau_b(m["predict"](test), [human_rating(r) for r in test]))
            btaus.append(kendall_tau_b([base.pointwise(r) for r in test], [human_rating(r) for r in test]))
            wtaus.append(kendall_tau_b([taste.pointwise(r) for r in test], [human_rating(r) for r in test]))
        rows.append({"n": n, "induced": mean(taus), "baseline": mean(btaus), "taste": mean(wtaus)})
    return rows


# ---------------------------------------------------------------- part 10: a small eval harness


class TransientError(Exception):
    """An error worth retrying: a timeout, a rate limit."""


class FlakyBackend:
    """Wraps a function so that some calls raise TransientError. Whether an attempt fails depends only on the key
    and the attempt number, so runs are repeatable."""

    def __init__(self, fn, fail_rate=0.2, seed=0):
        self.fn, self.fail_rate, self.seed = fn, fail_rate, seed
        self.calls = 0
        self.attempts = defaultdict(int)

    def __call__(self, task, rng):
        self.calls += 1
        # identify the call by the task and the generator state, without drawing from the generator
        key = task["id"] + "|" + hashlib.md5(repr(rng.getstate()).encode()).hexdigest()
        a = self.attempts[key] = self.attempts[key] + 1
        if random.Random(f"flaky|{self.seed}|{task['id']}|{a}").random() < self.fail_rate:
            raise TransientError("rate limit")
        return self.fn(task, rng)


class EvalHarness:
    """Runs `fn(task, rng)` for every (task, trial). Features:
    - a seed per (task, trial), so results do not depend on the order of the runs
    - a cache keyed by a hash of (version, config, seed, task, trial); only successes are stored
    - retries on TransientError with exponential backoff (the sleep function is injected, so tests do not wait)
    - a result log, one dict per attempt (JSON lines if log_path is given)
    - a cost budget: once spent, remaining runs are marked skipped
    `version` must change when the judge prompt, the model or the rubric changes, or old results will be reused."""

    def __init__(self, fn, version="v1", config=None, cache=None, max_retries=3, backoff=0.5, sleep=None,
                 seed=0, log_path=None, budget=None, cost_fn=None):
        self.fn, self.version, self.config = fn, version, dict(config or {})
        self.cache = cache if cache is not None else {}
        self.max_retries, self.backoff = max_retries, backoff
        self.sleep = sleep or (lambda s: None)
        self.seed, self.log_path, self.budget = seed, log_path, budget
        self.cost_fn = cost_fn or (lambda task, result: 1.0)
        self.log = []
        self.slept = 0.0
        self.spent = 0.0
        self._clock = 0

    def key(self, task, trial):
        blob = json.dumps({"v": self.version, "c": self.config, "s": self.seed, "t": task["id"],
                           "x": task.get("input"), "k": trial}, sort_keys=True, default=str)
        return hashlib.sha256(blob.encode()).hexdigest()[:16]

    def _record(self, **kw):
        self._clock += 1
        kw["step"] = self._clock
        self.log.append(kw)
        if self.log_path:
            with open(self.log_path, "a") as f:
                f.write(json.dumps(kw, sort_keys=True, default=str) + "\n")

    def run_one(self, task, trial=0):
        key = self.key(task, trial)
        base = {"task": task["id"], "trial": trial, "key": key, "version": self.version}
        if key in self.cache:
            self._record(status="cached", attempt=0, **base)
            return {"task": task["id"], "trial": trial, "status": "cached", "result": self.cache[key]}
        if self.budget is not None and self.spent >= self.budget:
            self._record(status="skipped", attempt=0, **base)
            return {"task": task["id"], "trial": trial, "status": "skipped", "result": None}
        for attempt in range(1, self.max_retries + 2):
            rng = random.Random(f"{self.seed}-{task['id']}-{trial}")
            try:
                res = self.fn(task, rng)
            except TransientError as e:
                self._record(status="retry" if attempt <= self.max_retries else "error", attempt=attempt,
                             error=str(e), **base)
                if attempt <= self.max_retries:
                    wait = self.backoff * 2 ** (attempt - 1)
                    self.slept += wait
                    self.sleep(wait)
                continue
            self.cache[key] = res
            self.spent += self.cost_fn(task, res)
            self._record(status="ok", attempt=attempt, **base)
            return {"task": task["id"], "trial": trial, "status": "ok", "result": res}
        return {"task": task["id"], "trial": trial, "status": "error", "result": None}

    def run(self, tasks, trials=1):
        return [self.run_one(t, k) for t in tasks for k in range(trials)]

    def summary(self):
        c = Counter(r["status"] for r in self.log)
        return {"ok": c["ok"], "cached": c["cached"], "retries": c["retry"], "errors": c["error"],
                "skipped": c["skipped"], "slept": self.slept, "spent": self.spent}


# ---------------------------------------------------------------- part 11: the demo


def _f(x, d=3):
    return f"{x:.{d}f}"


def standard_judges():
    """A panel of three judges, each with a different taste, and one judge for each single bias."""
    return {
        "fair": MockJudge("fair", seed=1),
        "first": MockJudge("first", pos_bias=0.40, seed=2),
        "wordy": MockJudge("wordy", len_bias=0.25, seed=3),
        "selfish": MockJudge("selfish", self_bias=0.35, self_system="epsilon", seed=4),
    }


def main():
    W = make_world(0)
    R = W["responses"]
    pairs = pair_list(W)
    J = standard_judges()
    print("== 1. code-based graders on the toy answers (the first of three grader types)")
    golds = {q["qid"]: q["gold"] for q in W["questions"]}
    right = [number_close(r["text"].split("The answer is ")[1], golds[r["qid"]]) for r in R]
    agree = percent_agreement(right, [r["traits"]["correct"] for r in R])
    print(f"number_close finds the right answer in {sum(right)} of {len(R)} responses; it agrees with the hidden trait on {_f(agree)}")
    hp = [human_pass(r) for r in R]
    print(f"but a human passes {sum(hp)} of {len(R)}; code-based pass (correct only) vs human pass: agreement {_f(percent_agreement(right, hp))}, kappa {_f(cohen_kappa(right, hp))}")
    print(f"\n== 2. position bias ({len(pairs)} pairs of different systems, each asked in both orders)")
    for nm in ("fair", "first"):
        r = position_report(J[nm], pairs)
        print(f"{nm:8s} consistency {_f(r['consistency'])}  first-slot rate {_f(r['first_rate'])}")
    print("\n== 3. verbosity bias (a response against itself plus 4 filler sentences)")
    sample = R[::6]
    for nm in ("fair", "wordy"):
        print(f"{nm:8s} padded copy gets {_f(verbosity_report(J[nm], W, sample)['padded_share'])} of the credit (0.5 is fair)")
    print("\n== 4. self-preference (the judge's own system is epsilon, which is weak and wordy)")
    neutral = MockJudge("selfish", self_bias=0.0, self_system="epsilon", seed=4)  # same judge, bias switched off
    for nm, jd in (("no bias", neutral), ("selfish", J["selfish"])):
        r = self_preference_report(jd, W)
        print(f"{nm:8s} judge win rate {_f(r['judge_winrate'])}  human win rate {_f(r['human_winrate'])}  gap {r['gap']:+.3f}")
    print("\n== 5. mitigations: swap, length control, a panel")
    fair, first, wordy = J["fair"], J["first"], J["wordy"]
    single = position_report(first, pairs)["consistency"]
    print(f"first-biased judge: consistency {_f(single)}; after swapping, verdicts are consistent by construction, and a disagreement becomes a tie")
    truth = [("a" if a["h"] > b["h"] else "b") for a, b in pairs]
    raw = [{"A": "a", "B": "b", "tie": "tie"}[first.pairwise(a, b)] for a, b in pairs]
    sw = [swapped_verdict(first, a, b) for a, b in pairs]
    cnt = lambda v: (sum(x == t for x, t in zip(v, truth)), sum(x not in (t, "tie") for x, t in zip(v, truth)), sum(x == "tie" for x in v))
    print(f"verdicts out of {len(pairs)} pairs against the hidden quality (right, wrong, tie):")
    print(f"  first-biased judge, raw order  : {cnt(raw)}")
    print(f"  first-biased judge, swapped    : {cnt(sw)}")
    rows = system_pair_rows(W, wordy, "epsilon", "delta")
    lc = length_controlled_winrate(rows)
    hum = mean(1.0 if W["by_q"][q]["epsilon"]["h"] > W["by_q"][q]["delta"]["h"] else 0.0 for q in W["by_q"])
    print(f"epsilon vs delta with a wordy judge: raw win rate {_f(lc['raw'])}, length-controlled {_f(lc['controlled'])}, human (from h) {_f(hum)}")
    panel = [J["fair"], J["first"], J["wordy"], J["selfish"]]
    for nm in ("fair", "first", "wordy", "selfish"):
        v = [swapped_verdict(J[nm], a, b) for a, b in pairs]
        print(f"  {nm:8s} alone, swapped: {cnt(v)}")
    pv = [panel_verdict(panel, a, b) for a, b in pairs]
    print(f"  panel of 4, swapped, no own-system votes: {cnt(pv)}")
    print("\n== 6. pointwise against pairwise: five systems ranked")
    tq = true_system_quality(W)
    print("true quality (mean h):", {s: round(v, 3) for s, v in tq.items()})
    pw = pointwise_system_scores(W, fair)
    comps = comparisons_from(W, lambda a, b: swapped_verdict(fair, a, b))
    bt = bradley_terry(comps, W["systems"])
    elo = to_elo(bt)
    print("pointwise mean score:", {s: round(v, 2) for s, v in pw.items()})
    print("Bradley-Terry Elo   :", {s: round(v) for s, v in elo.items()})
    print(f"judge calls: pointwise {len(R)}, pairwise {len(pairs)} pairs x 2 orders = {2 * len(pairs)}")
    print("agreement with truth: pointwise", {k: round(v, 3) for k, v in ranking_agreement(pw, tq).items()}, "BT", {k: round(v, 3) for k, v in ranking_agreement(bt, tq).items()})
    pw_w = pointwise_system_scores(W, wordy)
    bt_w = bradley_terry(comparisons_from(W, lambda a, b: swapped_verdict(wordy, a, b)), W["systems"])
    print("wordy judge, rank order: pointwise", sorted(pw_w, key=lambda s: -pw_w[s]), "BT", sorted(bt_w, key=lambda s: -bt_w[s]), "truth", sorted(tq, key=lambda s: -tq[s]))
    elos = [elo_online(comps, seed=i) for i in range(5)]
    print("online Elo, 5 shuffles, alpha's rating:", [round(e["alpha"]) for e in elos])
    boot = bootstrap_bt(comps, W["systems"], n_boot=60)
    print("bootstrap 95% interval of log-strength:", {s: (round(lo, 2), round(hi, 2)) for s, (lo, hi) in boot["ci"].items()})
    print("\n== 7. judge and human agreement")
    hl = [human_pass(r) for r in R]
    jl = [fair.pass_label(r) for r in R]
    print(f"fair judge pass/fail: agreement {_f(percent_agreement(jl, hl))}, kappa {_f(cohen_kappa(jl, hl))}, fail TPR/TNR {tpr_tnr(jl, hl)}")
    jr = [fair.pointwise(r) for r in R]
    hr = [human_rating(r) for r in R]
    print(f"fair judge 1 to 5: exact agreement {_f(percent_agreement(jr, hr))}, weighted kappa {_f(weighted_kappa(jr, hr))}, Spearman {_f(spearman(jr, hr))}")
    thr = imbalanced_labels(W, 0.05)
    hl2 = [r["h"] >= thr for r in R]
    jl2 = [True] * len(R)
    print(f"rare failures (5%): a judge that always says pass has agreement {_f(percent_agreement(jl2, hl2))}, kappa {_f(cohen_kappa(jl2, hl2))}, fail TPR {tpr_tnr(jl2, hl2)['tpr']:.2f}")
    confs = [fair.pass_confidence(r) for r in R]
    print(f"calibration: expected calibration error {_f(expected_calibration_error(confs, hl))}")
    a, b = platt_fit([c for c in confs[::2]], hl[::2], iters=3000, lr=2.0)
    cal = [sigmoid(a + b * c) for c in confs[1::2]]
    print(f"after a Platt fit on half: error on the other half {_f(expected_calibration_error(cal, hl[1::2]))} (before, same half: {_f(expected_calibration_error(confs[1::2], hl[1::2]))})")
    print("score to human pass rate:", {s: (round(p, 2), n) for s, (p, n) in score_to_pass_rate(jr, hl).items()})
    print("\n== 8. pass@k and pass^k with a noisy grader (the agent is the same in both rows)")
    rng = random.Random(5)
    task_p = [rng.choice([0.95, 0.9, 0.8, 0.6, 0.3]) for _ in range(60)]
    t, j = graded_runs(task_p, 16, tpr=0.92, tnr=0.90, seed=1)
    ks = [1, 2, 4, 8, 16]
    ct, cj = curves_from(t, ks), curves_from(j, ks)
    print("k            ", ks)
    print("true pass^k  ", [round(v, 3) for v in ct["pass_hat_k"]])
    print("judged pass^k", [round(v, 3) for v in cj["pass_hat_k"]])
    print("true pass@k  ", [round(v, 3) for v in ct["pass_at_k"]])
    print("judged pass@k", [round(v, 3) for v in cj["pass_at_k"]])
    rr = judge_rerun_agreement(fair, R[:120])
    print(f"the fair judge asked 5 times on 120 responses: all runs equal on {_f(rr['all_same'])}, mean flip rate {_f(rr['flip_rate'])}")
    print("\n== 9. an AutoMetrics-style metric induction toy")
    for row in induction_curve(W):
        print(f"labels {row['n']:4d}  induced metric tau {row['induced']:.3f}   holistic judge {row['baseline']:.3f}   holistic judge with its own taste {row['taste']:.3f}")
    tr = R[::2][:80]
    m = induce_metric(tr, [human_rating(r) for r in tr], n_keep=4)
    print("criteria kept from 80 labels:", {n: round(w, 2) for n, w in zip(m["names"], m["weights"])})
    print("\n== 10. the harness: seeds, cache, retries, log")
    tasks = [{"id": f"t{i}", "input": i} for i in range(30)]
    flaky = FlakyBackend(lambda task, rng: round(rng.random() + task["input"], 6), fail_rate=0.25, seed=3)
    h = EvalHarness(flaky, version="judge-v1", seed=7)
    first_pass = h.run(tasks, trials=2)
    s1 = h.summary()
    again = h.run(tasks, trials=2)
    s2 = h.summary()
    print("first run :", {k: v for k, v in s1.items() if k in ("ok", "cached", "retries", "errors", "slept")})
    print("second run:", {k: v for k, v in s2.items() if k in ("ok", "cached", "retries", "errors", "slept")}, "(all hits, no new backend calls)")
    h2 = EvalHarness(flaky, version="judge-v2", cache=h.cache, seed=7)
    h2.run(tasks[:5], trials=1)
    print("after changing the version, 5 tasks:", {k: v for k, v in h2.summary().items() if k in ("ok", "cached")})
    same = [r["result"] for r in first_pass if r["status"] in ("ok", "cached")] == [r["result"] for r in again if r["status"] in ("ok", "cached")]
    print("results identical on rerun:", same)
    hb = EvalHarness(lambda task, rng: 1, budget=10.0, seed=0)
    hb.run(tasks, 1)
    print("with a budget of 10 units at 1 per run:", {k: v for k, v in hb.summary().items() if k in ("ok", "skipped", "spent")})


if __name__ == "__main__":
    main()