"""Check the fast bootstrap against the literal one, and D_MH against a hand sum.

    python research/rain-before-seven-replication/test_bootstrap.py

The fast path sums precomputed per-block terms instead of rebuilding 76,000
rows per draw. That is a claim about algebra, so it gets checked on the real
data with the same seed: the two paths must return the same interval to well
under a hundredth of a point.
"""
import sys

import numpy as np

import load
import rule


def main():
    d = load.days(rule.SITES)
    R, F = d["R"], d["F"]
    bad = 0

    # 1. D_MH against a plain Python loop over blocks
    num = den = 0.0
    for b in range(d["n_blocks"]):
        m = d["block"] == b
        r, f = R[m], F[m]
        n1, n0 = int(r.sum()), int((~r).sum())
        if n1 == 0 or n0 == 0:
            continue
        w = n1 * n0 / (n1 + n0)
        num += w * (f[r].mean() - f[~r].mean())
        den += w
    ref = 100.0 * num / den
    got = load.d_mh(d["block"], d["n_blocks"], R, F)
    print(f"D_MH  loop {ref:+.9f}   vectorised {got:+.9f}")
    if abs(ref - got) > 1e-9:
        print("  FAIL"); bad += 1

    # 2. the two bootstrap paths, same seed, same reps
    for unit in ("ym", "block"):
        a = load.bootstrap(d, R, F, unit, reps=300, seed=7)
        b = load.bootstrap_slow(d, R, F, unit, reps=300, seed=7)
        print(f"boot {unit:<6} fast [{a[0]:+.6f}, {a[1]:+.6f}]   "
              f"slow [{b[0]:+.6f}, {b[1]:+.6f}]")
        if max(abs(a[0] - b[0]), abs(a[1] - b[1])) > 1e-6:
            print("  FAIL"); bad += 1

    # 3. an injection of a known size must land near that size
    from run import inject
    rng = np.random.default_rng(1)
    for pts in (2.0, 10.0):
        Fi = inject(d, R, F, pts, rng)
        v = load.d_mh(d["block"], d["n_blocks"], R, Fi)
        base = load.d_mh(d["block"], d["n_blocks"], R, F)
        print(f"inject {pts:>4.0f} pts -> shift {v - base:+.3f}")
        if abs((v - base) + pts) > 0.35:
            print("  FAIL"); bad += 1

    print("\n" + ("all checks pass" if bad == 0 else f"{bad} CHECK(S) FAILED"))
    return 1 if bad else 0


if __name__ == "__main__":
    sys.exit(main())