"""The two controls rule.py requires, run before the headline is believed.

    python research/rain-before-seven/controls.py

NULL: shuffle the R flag within each (site, month) block. |delta| must be
      under 2 points and its CI must contain 0.
POSITIVE: from the shuffled labels, flip F to 0 on a random subset of the
      R-labelled days so that the injected effect is exactly -5 points, the
      materiality bar. The estimator must recover -5 inside its CI and the CI
      must exclude 0 — otherwise an effect of the size the kill rule cares
      about is undetectable at this n, and the study is void: unfailable.
"""
import sys

import numpy as np

import load

SEED = 20260919


def shuffled_R(d, rng):
    """R permuted within each (site, month) block: same counts, no information."""
    R = d["R"].copy()
    for b in range(d["n_blocks"]):
        m = np.flatnonzero(d["block"] == b)
        R[m] = rng.permutation(R[m])
    return R


def main():
    d = load.days()
    rng = np.random.default_rng(SEED)
    print(f"days: {len(d['R'])}   blocks: {d['n_blocks']}")

    Rn = shuffled_R(d, rng)
    dn = load.delta(Rn, d["F"])
    lo, hi = load.block_bootstrap(d, Rn, d["F"], reps=2000, seed=SEED + 1)
    null_ok = abs(dn) < 2.0 and lo < 0 < hi
    print(f"\nNULL pooled  delta = {dn:+.3f} pts   CI [{lo:+.3f}, {hi:+.3f}]   "
          f"{'PASS' if null_ok else 'FAIL'}")

    dns = load.delta_strat(d["block"], d["n_blocks"], Rn, d["F"])
    slo, shi = load.block_bootstrap(d, Rn, d["F"], reps=2000, seed=SEED + 1,
                                    stat="strat")
    null_s_ok = abs(dns) < 2.0 and slo < 0 < shi
    print(f"NULL strat   delta = {dns:+.3f} pts   CI [{slo:+.3f}, {shi:+.3f}]   "
          f"{'PASS' if null_s_ok else 'FAIL'}")

    # Inject exactly -5 points: among the shuffled R days, flip F 1->0 on
    # k = round(0.05 * n_R) of the days where F is currently 1.
    F = d["F"].copy()
    ridx = np.flatnonzero(Rn)
    k = int(round(0.05 * len(ridx)))
    ones = ridx[F[ridx]]
    flip = rng.choice(ones, size=k, replace=False)
    F[flip] = False
    dp = load.delta(Rn, F)
    plo, phi = load.block_bootstrap(d, Rn, F, reps=2000, seed=SEED + 2)
    pos_ok = plo <= -5.0 <= phi and phi < 0
    print(f"\nPOSITIVE pooled  injected -5.000, measured {dp:+.3f} pts   "
          f"CI [{plo:+.3f}, {phi:+.3f}]   {'PASS' if pos_ok else 'FAIL'}")

    dps = load.delta_strat(d["block"], d["n_blocks"], Rn, F)
    pslo, pshi = load.block_bootstrap(d, Rn, F, reps=2000, seed=SEED + 2,
                                      stat="strat")
    pos_s_ok = pslo <= -5.0 <= pshi and pshi < 0
    print(f"POSITIVE strat   injected -5.000, measured {dps:+.3f} pts   "
          f"CI [{pslo:+.3f}, {pshi:+.3f}]   {'PASS' if pos_s_ok else 'FAIL'}")

    both = null_ok and pos_ok
    print("\n" + ("SEALED STATISTIC: both controls pass."
                  if both else
                  "SEALED STATISTIC: a control FAILED — rule.py forbids a headline."))
    print("STRATIFIED FOLLOW-UP: "
          + ("both controls pass." if (null_s_ok and pos_s_ok) else "a control FAILED."))
    return 0 if both else 1


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