"""Score the sealed rule `rule.rain-before-seven-replication`.

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

The gates run and print FIRST, before the real D_MH is computed, because
rule.py says so. If a gate fails, nothing is scored and the study is a void.

Control intervals use fewer bootstrap draws than the headline (the seal fixes
BOOTSTRAPS = 10,000 for the interval that carries the verdict and says nothing
about the 200-repetition controls, where 200 x 10,000 is not computable in a
two-hour session). The number used is printed with the result.
"""
import argparse
import json
import sys
import time

import numpy as np

import load
import rule

SEED = rule.SEED


def shuffled_R(d, rng):
    """R permuted within each verdict block: same counts, no information."""
    R = d["R"].copy()
    for rows in load._members(d["block"], d["n_blocks"]):
        R[rows] = rng.permutation(R[rows])
    return R


def inject(d, R, F, points, rng):
    """Flip F 1->0 on R days so the within-block difference falls by `points`.

    Per block b the flip probability is points/100 divided by that block's
    P(F | R); a block whose P(F|R) is already below the target has every one of
    its R-and-F days flipped, which is the most that can be injected there.
    Done per block rather than with one global rate because a single rate
    injects a different number of points in every block.
    """
    F = F.copy()
    for rows in load._members(d["block"], d["n_blocks"]):
        r = rows[R[rows]]
        if len(r) == 0:
            continue
        ones = r[F[r]]
        p1 = len(ones) / len(r)
        if p1 == 0:
            continue
        p = min(1.0, (points / 100.0) / p1)
        # each candidate day flipped independently with probability p, rather
        # than a rounded count: with ~8 R-and-F days in a block, rounding
        # p*len(ones) to a whole day sends every injection under ~6 points to
        # zero in almost every block, and the floor G3 measures would be an
        # artefact of that rounding. Caught by test_bootstrap.py.
        F[ones[rng.random(len(ones)) < p]] = False
    return F


def excludes_zero(lo, hi):
    return (lo > 0) or (hi < 0)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--control-reps", type=int, default=400)
    ap.add_argument("--control-runs", type=int, default=200)
    ap.add_argument("--reps", type=int, default=rule.BOOTSTRAPS)
    ap.add_argument("--out", default="results.json")
    a = ap.parse_args()

    t0 = time.time()
    d = load.days(rule.SITES)
    R, F = d["R"], d["F"]
    n = len(R)
    res = {"n_site_days": int(n), "dropped_days": int(d["dropped"]),
           "n_blocks": int(d["n_blocks"]), "n_blocks96": int(d["n_blocks96"]),
           "n_ym": int(d["n_ym"]), "control_reps": a.control_reps,
           "control_runs": a.control_runs, "reps": a.reps}
    print(f"site-days {n}   dropped {d['dropped']}   "
          f"blocks(site,ym) {d['n_blocks']}   blocks96 {d['n_blocks96']}   "
          f"months {d['n_ym']}")
    print(f"P(R) = {R.mean():.4f}")

    # ---- G4 SANITY --------------------------------------------------------
    pF = float(F.mean())
    g4 = 0.60 <= pF <= 0.95
    res["G4"] = {"P_F": pF, "pass": bool(g4)}
    print(f"\nG4 SANITY   P(F) = {pF:.4f}   "
          f"must lie in [0.60, 0.95]   {'PASS' if g4 else 'FAIL'}")
    if not g4:
        print("VOID: the fetch or the clock is wrong. Nothing is scored.")
        json.dump(res, open(a.out, "w"), indent=1)
        return 1

    # ---- G1 NULL ----------------------------------------------------------
    rng = np.random.default_rng(SEED + 1)
    ds, excl = [], 0
    for i in range(a.control_runs):
        Rn = shuffled_R(d, rng)
        ds.append(load.d_mh(d["block"], d["n_blocks"], Rn, F))
        lo, hi = load.bootstrap(d, Rn, F, "ym", reps=a.control_reps,
                                seed=SEED + 1000 + i)
        excl += excludes_zero(lo, hi)
        if i % 25 == 0:
            print(f"   G1 {i}/{a.control_runs}  ({time.time()-t0:.0f}s)",
                  flush=True)
    ds = np.array(ds)
    frac = excl / a.control_runs
    g1 = abs(ds.mean()) < 1.0 and frac < 0.12
    res["G1"] = {"mean": float(ds.mean()), "sd": float(ds.std()),
                 "excl_frac": frac, "pass": bool(g1)}
    print(f"\nG1 NULL     mean D_MH over {a.control_runs} shuffles = "
          f"{ds.mean():+.4f} pts (sd {ds.std():.4f}); "
          f"interval excluded 0 on {frac:.1%}")
    print(f"            needs |mean| < 1.0 and share < 12 %   "
          f"{'PASS' if g1 else 'FAIL'}")

    # ---- G2 POSITIVE ------------------------------------------------------
    rng = np.random.default_rng(SEED + 2)
    got, excl = [], 0
    for i in range(a.control_runs):
        Fi = inject(d, R, F, 10.0, rng)
        got.append(load.d_mh(d["block"], d["n_blocks"], R, Fi))
        lo, hi = load.bootstrap(d, R, Fi, "ym", reps=a.control_reps,
                                seed=SEED + 2000 + i)
        excl += excludes_zero(lo, hi)
        if i % 25 == 0:
            print(f"   G2 {i}/{a.control_runs}  ({time.time()-t0:.0f}s)",
                  flush=True)
    got = np.array(got)
    rate = excl / a.control_runs
    g2 = rate >= 0.90
    res["G2"] = {"measured_mean": float(got.mean()), "excl_frac": rate,
                 "pass": bool(g2)}
    print(f"\nG2 POSITIVE injected -10 pts on the real rows; measured "
          f"{got.mean():+.3f} pts; interval excluded 0 on {rate:.1%}")
    print(f"            needs >= 90 %   {'PASS' if g2 else 'FAIL'}")

    if not (g1 and g2):
        print("\nVOID: a gate failed. rule.py forbids any verdict from this "
              "estimator, and forbids repairing it inside this run.")
        json.dump(res, open(a.out, "w"), indent=1)
        return 1

    # ---- G3 FLOOR ---------------------------------------------------------
    floor, g3_detail = None, {}
    for pts in (1.0, 2.0, 5.0, 10.0):
        rng = np.random.default_rng(SEED + 3)
        excl = 0
        for i in range(a.control_runs):
            Fi = inject(d, R, F, pts, rng)
            lo, hi = load.bootstrap(d, R, Fi, "ym", reps=a.control_reps,
                                    seed=SEED + 3000 + i)
            excl += excludes_zero(lo, hi)
        r = excl / a.control_runs
        g3_detail[str(pts)] = r
        print(f"\nG3 FLOOR    {pts:>4.0f} pts recovered on {r:.1%} of "
              f"{a.control_runs} injections   ({time.time()-t0:.0f}s)",
              flush=True)
        if r >= 0.90 and floor is None:
            floor = pts
            break
    res["G3"] = {"floor": floor, "detail": g3_detail,
                 "unfailable": floor is None or floor > 5.0}
    print(f"\nG3 FLOOR    smallest effect recovered >= 90 %: "
          f"{floor if floor is not None else '>10'} pts")
    if res["G3"]["unfailable"]:
        print("            floor above C1's bar of 5 pts: "
              "C1 is `void: unfailable`, and that is said before the number.")

    # ---- THE VERDICT ------------------------------------------------------
    D = load.d_mh(d["block"], d["n_blocks"], R, F)
    lo, hi = load.bootstrap(d, R, F, "ym", reps=a.reps, seed=SEED)
    slo, shi = load.bootstrap(d, R, F, "block", reps=a.reps, seed=SEED)
    D96 = load.d_mh(d["block96"], d["n_blocks96"], R, F)
    res["D_MH"] = D
    res["primary_ci"] = [lo, hi]
    res["secondary_ci"] = [slo, shi]
    res["D_MH_block96"] = D96
    print(f"\n=== D_MH = {D:+.3f} points")
    print(f"    PRIMARY   (months, all sites at once)  95 % [{lo:+.3f}, {hi:+.3f}]"
          f"   width {hi-lo:.3f}")
    print(f"    SECONDARY (site-month blocks)          95 % [{slo:+.3f}, {shi:+.3f}]"
          f"   width {shi-slo:.3f}")
    wider = (hi - lo) > (shi - slo)
    res["primary_wider"] = bool(wider)
    print(f"    the seal predicted the primary would be WIDER: "
          f"{'yes' if wider else 'NO -- something is wrong with one of them'}")
    print(f"    sensitivity, 96-block reading of the seal: D_MH = {D96:+.3f}")

    if res["G3"]["unfailable"]:
        c1 = "void: unfailable"
    elif hi < rule.BAR:
        c1 = "KILLED"
    elif lo > -rule.BAR:
        c1 = "SURVIVED"
    else:
        c1 = "INCONCLUSIVE"
    c2 = "REPLICATED" if lo <= rule.PUBLISHED <= hi else "NOT REPLICATED"
    res["C1"], res["C2"] = c1, c2
    print(f"\nC1 (proverb, bar {rule.BAR:+.0f}): {c1}")
    print(f"C2 (replication of {rule.PUBLISHED:+.2f}): {c2}")

    # ---- SECONDARIES ------------------------------------------------------
    Ds = load.d_mh(d["block"], d["n_blocks"], R, d["Fstrict"])
    slo2, shi2 = load.bootstrap(d, R, d["Fstrict"], "ym", reps=a.reps,
                                seed=SEED)
    res["S1"] = {"D_MH": Ds, "ci": [slo2, shi2],
                 "P_Fstrict": float(d["Fstrict"].mean())}
    print(f"\nS1 F_strict (08:00-10:59 all dry): D_MH = {Ds:+.3f}   "
          f"95 % [{slo2:+.3f}, {shi2:+.3f}]   P(F_strict) = "
          f"{d['Fstrict'].mean():.4f}")

    print("\nS2 per site:")
    res["S2_site"] = {}
    for si, site in enumerate(d["sites"]):
        m = d["site"] == si
        b = d["block"][m]
        _, bi = np.unique(b, return_inverse=True)
        v = load.d_mh(bi, int(bi.max()) + 1, R[m], F[m])
        res["S2_site"][site] = v
        print(f"   {site:<10} n={m.sum():>6}  D_MH = {v:+.3f}")

    print("\nS2 per calendar month:")
    res["S2_month"] = {}
    mon = np.array([x[5:7] for x in d["date"]])
    for mm in sorted(set(mon)):
        m = mon == mm
        b = d["block"][m]
        _, bi = np.unique(b, return_inverse=True)
        v = load.d_mh(bi, int(bi.max()) + 1, R[m], F[m])
        res["S2_month"][mm] = v
        print(f"   {mm}  n={m.sum():>6}  D_MH = {v:+.3f}")

    # S3: the same pipeline in the Pacific Northwest. NOT a test of the proverb.
    dp = load.days(rule.PNW_SITES)
    Dp = load.d_mh(dp["block"], dp["n_blocks"], dp["R"], dp["F"])
    plo, phi = load.bootstrap(dp, dp["R"], dp["F"], "ym", reps=a.reps,
                              seed=SEED)
    res["S3"] = {"D_MH": Dp, "ci": [plo, phi], "n": int(len(dp["R"])),
                 "P_R": float(dp["R"].mean()), "P_F": float(dp["F"].mean())}
    print(f"\nS3 Pacific Northwest (NOT a test of the English proverb): "
          f"n={len(dp['R'])}  P(R)={dp['R'].mean():.4f}  "
          f"P(F)={dp['F'].mean():.4f}")
    print(f"   D_MH = {Dp:+.3f}   95 % [{plo:+.3f}, {phi:+.3f}]")
    res["S3_site"] = {}
    for si, site in enumerate(dp["sites"]):
        m = dp["site"] == si
        b = dp["block"][m]
        _, bi = np.unique(b, return_inverse=True)
        v = load.d_mh(bi, int(bi.max()) + 1, dp["R"][m], dp["F"][m])
        res["S3_site"][site] = v
        print(f"   {site:<10} n={m.sum():>6}  D_MH = {v:+.3f}")

    json.dump(res, open(a.out, "w"), indent=1, default=float)
    print(f"\nwrote {a.out} in {time.time()-t0:.0f}s")
    return 0


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