"""The statistic and its interval, exactly as `rule.py` defines them.

Shared by `controls.py` and `analyse.py` so that the controls and the real
result cannot drift apart — a control that exercises a different code path
from the thing it is controlling is testing the paperwork
(`memory/controls-test-the-detector-not-the-paperwork.md`). `rule.py` names
five files and this is a sixth; it adds no definition, it holds the one both
of the others use.

THE STATISTIC. delta = a weighted mean, over strata, of the within-stratum
difference P(outcome | label) - P(outcome | not label), in percentage points.
A stratum is (site, calendar month of the year): 14 x 12 = 168 of them. The
weight is the stratum's site-day count. Strata with no exposed day, or no
unexposed day, contribute nothing and are counted.

THE BLOCK. A bootstrap block is (site, calendar month), 14 x 120 = 1,680 of
them. Every block lies inside exactly one stratum, which is what makes the
resampling cheap: a block carries four numbers (n1, s1, n0, s0) and a stratum
is the sum of its ten blocks. The resample draws 1,680 blocks with replacement
and recomputes the whole statistic from those sums.

That block choice asserts independence ACROSS SITES, which is false — Cardiff
and Swansea share their weather — so the interval is narrower than the truth
by a factor nobody here has measured. `rule.py` says so, and a verdict that
turns on the last tenth of the interval has to say so too.
"""
import numpy as np

N_RESAMPLES = 10000
SEED = 20260923


def blocks(rows, label_key, outcome_key):
    """Per-block (n1, s1, n0, s0) and the stratum each block belongs to.

    Rows whose label or outcome is -1 (unknown) are dropped here, per arm,
    which is where `load.py`'s -1 convention is finally applied.
    """
    keys, strata = {}, {}
    n1 = []
    s1 = []
    n0 = []
    s0 = []
    stratum_of = []
    for r in rows:
        lab, out = r[label_key], r[outcome_key]
        if lab not in (0, 1) or out not in (0, 1):
            continue
        block = (r["site"], r["date"][:7])
        stratum = (r["site"], r["month"])
        if block not in keys:
            keys[block] = len(n1)
            n1.append(0)
            s1.append(0)
            n0.append(0)
            s0.append(0)
            stratum_of.append(strata.setdefault(stratum, len(strata)))
        i = keys[block]
        if lab == 1:
            n1[i] += 1
            s1[i] += out
        else:
            n0[i] += 1
            s0[i] += out
    return (np.array(n1, float), np.array(s1, float),
            np.array(n0, float), np.array(s0, float),
            np.array(stratum_of), len(strata))


def delta(n1, s1, n0, s0, stratum_of, n_strata, counts=None):
    """The stratified statistic in percentage points, from block aggregates.

    `counts` is an optional per-block multiplicity, which is how a bootstrap
    resample is expressed without rebuilding anything.
    """
    if counts is not None:
        n1, s1, n0, s0 = n1 * counts, s1 * counts, n0 * counts, s0 * counts
    N1 = np.bincount(stratum_of, n1, n_strata)
    S1 = np.bincount(stratum_of, s1, n_strata)
    N0 = np.bincount(stratum_of, n0, n_strata)
    S0 = np.bincount(stratum_of, s0, n_strata)
    live = (N1 > 0) & (N0 > 0)
    if not live.any():
        return float("nan")
    d = S1[live] / N1[live] - S0[live] / N0[live]
    w = N1[live] + N0[live]
    return 100.0 * float(np.sum(w * d) / np.sum(w))


def interval(packed, n_resamples=N_RESAMPLES, seed=SEED):
    """(point, lo, hi) — the 95 % percentile block bootstrap."""
    n1, s1, n0, s0, stratum_of, n_strata = packed
    point = delta(n1, s1, n0, s0, stratum_of, n_strata)
    rng = np.random.default_rng(seed)
    nb = len(n1)
    out = np.empty(n_resamples)
    for i in range(n_resamples):
        counts = np.bincount(rng.integers(0, nb, nb), minlength=nb).astype(float)
        out[i] = delta(n1, s1, n0, s0, stratum_of, n_strata, counts)
    lo, hi = np.nanpercentile(out, [2.5, 97.5])
    return point, float(lo), float(hi)


def verdict(lo, hi, bar, direction):
    """rule.py's kill rule. direction +1 for the night arm, -1 for the morning."""
    if direction > 0:
        if lo > bar:
            return "SURVIVED"
        if hi < bar:
            return "KILLED"
    else:
        if hi < -bar:
            return "SURVIVED"
        if lo > -bar:
            return "KILLED"
    return "INCONCLUSIVE"


def read_rows(path):
    import csv
    with open(path, newline="", encoding="utf-8") as fh:
        rows = []
        for r in csv.DictReader(fh):
            if r["dropped_reason"]:
                continue
            for k in ("month", "clear_w", "lit", "red_night", "clear_e",
                      "red_morning", "clear_e_at_set", "lit_rise",
                      "fine_next", "fine_today"):
                r[k] = int(r[k])
            rows.append(r)
    return rows
