#!/usr/bin/env python3
"""Confidence intervals that account for passages being clustered inside books.

Why this exists. A Wilson interval on 12,247 passages treats every passage as an independent
draw. They are not: they come from 1,809 books, and passages from one book share an author, a
register, a translator and one scanner's OCR quality. If a detector's verdict is really a
property of the *book*, the effective sample size is closer to the number of books than to the
number of passages, and a passage-level interval is too narrow — it claims a precision the design
does not have.

So the published interval is a **cluster bootstrap**: resample whole books with replacement,
pool their passages, recompute the rate. That is the standard fix and it makes no distributional
assumption. What it does *not* come with is a guarantee of being wider: clustering usually widens
an interval, but a percentile bootstrap is a different estimator from a Wilson interval and can
land inside it — on a subgroup with few books, or simply by sampling variation. An earlier version
of this docstring said it "can only ever be more conservative", which was the comfortable belief
rather than the fact, and it is the belief the caller was already built not to rely on.

Added after scoring began. Worth saying plainly, because a method added mid-study is normally a
red flag: it never moves a point estimate, and it is not chosen per row by which answer is nicer.
Both intervals are published side by side and the *wider* of the two is the headline, row by row,
which is the weaker claim in every direction — see the selection in build_page.py, which also
prints how many rows each method won.

Imported by analyse.py and build_page.py rather than reimplemented in each — the drift between
two hand-written copies of one rule always lands on the branch the check does not reach.
"""
import math
import random

SEED = 20260828          # the same seed the scan order uses; fixed, never re-drawn
B = 2000                 # bootstrap replicates


def wilson(k, n, z=1.96):
    """Passage-level interval. Wilson rather than normal: at single-digit rates a normal
    interval can dip below zero, which would be a nonsense number to publish."""
    if not n:
        return (0.0, 0.0)
    p = k / n
    d = 1 + z * z / n
    c = (p + z * z / (2 * n)) / d
    h = z * math.sqrt(p * (1 - p) / n + z * z / (4 * n * n)) / d
    return (max(0.0, c - h), min(1.0, c + h))


def cluster_bootstrap(clusters, b=B, seed=SEED):
    """clusters: list of (flagged, total) per source item. Returns (lo, hi) at 95%.

    Resamples the *books*, not the passages. A book contributes all of its scored passages or
    none of them, which is exactly the dependence a passage-level interval ignores.
    """
    clusters = [c for c in clusters if c[1]]
    m = len(clusters)
    if m < 2:
        return (0.0, 1.0)
    rng = random.Random(seed)
    ks = [c[0] for c in clusters]
    ns = [c[1] for c in clusters]
    reps = []
    for _ in range(b):
        K = N = 0
        for _ in range(m):
            i = rng.randrange(m)
            K += ks[i]
            N += ns[i]
        if N:
            reps.append(K / N)
    reps.sort()
    lo = reps[int(0.025 * (len(reps) - 1))]
    hi = reps[int(0.975 * (len(reps) - 1))]
    return (lo, hi)


def diff_bootstrap(clusters_a, clusters_b, b=B, seed=SEED):
    """95% CI on (rate_a - rate_b), resampling books independently within each corpus.

    Fixed before the control corpus was scored. The whole study turns on this one comparison, and
    picking a test after seeing both rates is how a difference becomes significant. Bootstrapping
    the difference directly — rather than eyeballing whether two intervals overlap — is also the
    correct thing to do: non-overlapping intervals imply a difference, but overlapping ones do
    not imply its absence, and that asymmetry always seems to be discovered in whichever
    direction suits the author.
    """
    A = [c for c in clusters_a if c[1]]
    B_ = [c for c in clusters_b if c[1]]
    if len(A) < 2 or len(B_) < 2:
        return (-1.0, 1.0)
    rng = random.Random(seed)
    reps = []
    for _ in range(b):
        out = []
        for cl in (A, B_):
            m = len(cl)
            K = N = 0
            for _ in range(m):
                i = rng.randrange(m)
                K += cl[i][0]
                N += cl[i][1]
            out.append(K / N if N else 0.0)
        reps.append(out[0] - out[1])
    reps.sort()
    return (reps[int(0.025 * (len(reps) - 1))], reps[int(0.975 * (len(reps) - 1))])


def by_cluster(rows, scores, idkey, thresh=0.5):
    """Collapse scored passages into one (flagged, total) pair per source item.

    Returned in book-id order, not in the order the rows arrived. That is not tidiness: the
    bootstrap samples cluster *positions* under a fixed seed, so the same books in a different
    order are a different sequence of draws and a different interval — in the second decimal,
    which is exactly the precision this study publishes to. And the two programs that compute
    these intervals walk the corpus in different orders by construction: `build_page.py` reads
    passages.jsonl in file order, `analyse.py` iterates the score file, which is in the scan's
    seeded-shuffle order. Same data, same seed, two different published bounds — with the page
    and analysis.txt each claiming to hold the other's numbers.

    Sorting here makes the interval a function of the data alone, which is what a number a reader
    is invited to recompute has to be. Same reason padding_check sorts before it shuffles.
    """
    agg = {}
    for r in rows:
        s = scores.get(r["n"])
        if not s:
            continue
        a = agg.setdefault(r[idkey], [0, 0])
        a[1] += 1
        a[0] += s["p_ai"] > thresh
    return [tuple(agg[k]) for k in sorted(agg, key=str)]


def both(rows, scores, idkey, thresh=0.5):
    """(k, n, n_clusters, wilson_lo, wilson_hi, boot_lo, boot_hi, design_effect_on_width)."""
    cl = by_cluster(rows, scores, idkey, thresh)
    k = sum(c[0] for c in cl)
    n = sum(c[1] for c in cl)
    wlo, whi = wilson(k, n)
    blo, bhi = cluster_bootstrap(cl)
    ww = whi - wlo
    return (k, n, len(cl), wlo, whi, blo, bhi, (bhi - blo) / ww if ww else float("nan"))