"""bootstrap-eval-deltas.py -- confidence intervals and a paired test for an eval delta on 500 items.

What it does: simulates two systems scored on the same 500-item test set with a known true
accuracy gap (the truth is a parameter so you can see whether the interval covers it), then
computes: (1) each system's accuracy with a percentile bootstrap 95% CI; (2) the PAIRED
bootstrap CI of the difference, resampling items so that both systems' scores on an item
stay together; (3) the unpaired CI of the difference, to show how much wider it is when
correlation between systems is ignored; (4) McNemar's exact test on the discordant pairs;
and (5) a coverage check: over 1,000 replays of the experiment, how often the paired CI
contains the true gap, and how often a "significant" result would have been declared.

Everything here is computed from a seeded simulation; the constants at the top are inputs.

Inputs:  none (edit N_ITEMS, ACC_A, ACC_B, CORR)
Run:     python code/bootstrap-eval-deltas.py
Output:  datasets/bootstrap-eval-deltas.csv (per-item scores of the reported run)
Requires: Python 3.11+ standard library only.
"""
from __future__ import annotations

import csv
import math
import random
from statistics import NormalDist

N_ITEMS = 500
ACC_A = 0.72          # true accuracy of system A
ACC_B = 0.76          # true accuracy of system B (true gap = +4 points)
CORR = 0.6            # how much item difficulty is shared between the systems
B = 2000              # bootstrap resamples


_ND = NormalDist()
_ZA, _ZB = _ND.inv_cdf(ACC_A), _ND.inv_cdf(ACC_B)


def simulate(rng: random.Random) -> list[tuple[int, int]]:
    """Per-item (a_correct, b_correct). A shared Gaussian item-difficulty term makes the two
    systems' scores correlated (a Gaussian copula), while each marginal hits its target accuracy."""
    out = []
    w = math.sqrt(1 - CORR ** 2)
    for _ in range(N_ITEMS):
        d = rng.gauss(0, 1)
        za = CORR * d + w * rng.gauss(0, 1)
        zb = CORR * d + w * rng.gauss(0, 1)
        out.append((int(za < _ZA), int(zb < _ZB)))
    return out


def pct_ci(samples: list[float], level: float = 0.95) -> tuple[float, float]:
    s = sorted(samples)
    lo = s[int((1 - level) / 2 * len(s))]
    hi = s[int((1 + level) / 2 * len(s)) - 1]
    return lo, hi


def paired_bootstrap(scores: list[tuple[int, int]], rng: random.Random) -> tuple[float, float]:
    n = len(scores)
    deltas = []
    for _ in range(B):
        idx = [rng.randrange(n) for _ in range(n)]
        deltas.append(sum(scores[i][1] - scores[i][0] for i in idx) / n)
    return pct_ci(deltas)


def unpaired_bootstrap(scores: list[tuple[int, int]], rng: random.Random) -> tuple[float, float]:
    n = len(scores)
    a = [s[0] for s in scores]
    b = [s[1] for s in scores]
    deltas = []
    for _ in range(B):
        ma = sum(a[rng.randrange(n)] for _ in range(n)) / n
        mb = sum(b[rng.randrange(n)] for _ in range(n)) / n
        deltas.append(mb - ma)
    return pct_ci(deltas)


def mcnemar_exact(scores: list[tuple[int, int]]) -> tuple[int, int, float]:
    b_only = sum(1 for a, b in scores if b and not a)
    a_only = sum(1 for a, b in scores if a and not b)
    n = a_only + b_only
    k = min(a_only, b_only)
    # two-sided exact binomial p-value at p=0.5
    p = sum(math.comb(n, i) for i in range(0, k + 1)) / 2 ** n * 2 if n else 1.0
    return a_only, b_only, min(1.0, p)


def main() -> None:
    rng = random.Random(2026)
    scores = simulate(rng)
    acc_a = sum(a for a, _ in scores) / N_ITEMS
    acc_b = sum(b for _, b in scores) / N_ITEMS
    n = N_ITEMS
    ci_a = pct_ci([sum(scores[rng.randrange(n)][0] for _ in range(n)) / n for _ in range(B)])
    ci_b = pct_ci([sum(scores[rng.randrange(n)][1] for _ in range(n)) / n for _ in range(B)])
    paired = paired_bootstrap(scores, rng)
    unpaired = unpaired_bootstrap(scores, rng)
    a_only, b_only, p = mcnemar_exact(scores)
    print(f"n={N_ITEMS}, true accuracies A={ACC_A:.2f} B={ACC_B:.2f}, true gap=+{100 * (ACC_B - ACC_A):.1f} pts")
    print(f"observed: A={acc_a:.3f} [{ci_a[0]:.3f}, {ci_a[1]:.3f}]   B={acc_b:.3f} [{ci_b[0]:.3f}, {ci_b[1]:.3f}]")
    print(f"observed gap: {100 * (acc_b - acc_a):+.1f} pts")
    print(f"  paired bootstrap 95% CI:   [{100 * paired[0]:+.1f}, {100 * paired[1]:+.1f}] pts")
    print(f"  unpaired bootstrap 95% CI: [{100 * unpaired[0]:+.1f}, {100 * unpaired[1]:+.1f}] pts")
    print(f"  discordant pairs: B-only correct={b_only}, A-only correct={a_only}; McNemar exact p={p:.4f}")

    # coverage and power over replays
    covered = signif = 0
    replays = 1000
    for _ in range(replays):
        s = simulate(rng)
        lo, hi = paired_bootstrap(s, random.Random(rng.random()))
        if lo <= ACC_B - ACC_A <= hi:
            covered += 1
        if lo > 0:
            signif += 1
    print(f"\nover {replays} replays: paired CI covers the true gap {covered / replays:.1%} of the time; "
          f"CI excludes zero (a 'significant' +gap) {signif / replays:.1%} of the time")

    with open("datasets/bootstrap-eval-deltas.csv", "w", newline="", encoding="utf-8") as f:
        w = csv.writer(f)
        w.writerow(["item", "system_a_correct", "system_b_correct"])
        for i, (a, b) in enumerate(scores, 1):
            w.writerow([i, a, b])


if __name__ == "__main__":
    main()