"""Run the sealed rule (rule.py) on the fetched files. Writes results.json."""
import json
import random
import sys
from pathlib import Path

HERE = Path(__file__).parent
sys.path.insert(0, str(HERE))
import rule  # noqa: E402  (this study's)
import importlib.util  # noqa: E402

spec = importlib.util.spec_from_file_location("wr", HERE.parent / "weekly-rain" / "rule.py")
wr = importlib.util.module_from_spec(spec); spec.loader.exec_module(wr)


def _load(path):
    import datetime as dt, gzip
    out = {}
    with gzip.open(path, "rt") as f:
        for line in f:
            p = line.rstrip("\n").split(",")
            if p[2] != "PRCP" or p[5] != "":
                continue
            d = dt.date(int(p[1][:4]), int(p[1][4:6]), int(p[1][6:8]))
            if 6 <= d.month <= 8:
                out[d] = int(p[3])
    return out


def table(all_data, y0, y1):
    by_year, dropped = wr.summers(all_data, y0, y1)
    t = {}
    for y, days in by_year.items():
        row = [[0, 0, 0] for _ in range(7)]
        for wd, v in days:
            row[wd][0] += v; row[wd][1] += 1; row[wd][2] += v >= 10
        t[y] = row
    return t, dropped


def parts(row, a, b, shift=0, wet=False, boost=1.0):
    m = nm = e = ne = 0.0
    k = 2 if wet else 0
    for wd in range(7):
        lab = (wd + shift) % 7
        if lab in a:
            m += row[wd][k] * boost; nm += row[wd][1]
        elif lab in b:
            e += row[wd][k]; ne += row[wd][1]
    return (m, nm, e, ne)


def arm(t, a, b, wet=False):
    yp = {y: parts(r, a, b, wet=wet) for y, r in t.items()}
    d = wr.D(list(yp.values()))
    lo, hi = wr.boot(yp, 4000, 20261008)
    return {"D": d, "lo": lo, "hi": hi, "verdict": rule.verdict(d, lo, hi), "summers": len(t)}


def daily(t):
    names = "Mon Tue Wed Thu Fri Sat Sun".split()
    return {names[w]: sum(t[y][w][0] for y in t) / sum(t[y][w][1] for y in t) / 10 for w in range(7)}


def controls(t, a, b, real_d):
    rng = random.Random(7)
    ys = sorted(t)
    excl = surv = ge = inj_surv = 0
    for i in range(2000):
        sh = {y: rng.randrange(7) for y in ys}
        yp = {y: parts(t[y], a, b, shift=sh[y]) for y in ys}
        d = wr.D(list(yp.values()))
        lo, hi = wr.boot(yp, 500, 1000 + i)
        excl += lo > 0 or hi < 0
        surv += rule.verdict(d, lo, hi) == "SURVIVES"
        ge += d >= real_d
        if i < 200:
            yb = {y: parts(t[y], a, b, shift=sh[y], boost=1.10) for y in ys}
            db = wr.D(list(yb.values()))
            lob, hib = wr.boot(yb, 500, 5000 + i)
            inj_surv += rule.verdict(db, lob, hib) == "SURVIVES"
    return {"draws": 2000, "excludes_zero": excl / 2000, "survives": surv / 2000,
            "perm_p_one_sided": ge / 2000, "inject_survives": inj_surv / 200}


def main():
    data = {s: _load(rule.DATA / f"{s}.csv.gz") for s in rule.STATIONS if (rule.DATA / f"{s}.csv.gz").exists()}
    res = {"stations_loaded": len(data)}
    t, dropped = table(data, *rule.PRIMARY)
    res["stations_with_20_summers"] = sum(
        1 for rec in data.values()
        if sum(1 for y in range(1998, 2026) if sum(1 for d in rec if d.year == y) >= 85) >= 20)
    res["dropped_station_summers"] = dropped
    res["station_days"] = sum(c[1] for r in t.values() for c in r)
    res["primary"] = arm(t, rule.WEEK, rule.WKND)
    res["A_wetday"] = arm(t, rule.WEEK, rule.WKND, wet=True)
    res["B_bell_split_new_stations"] = arm(t, rule.BELL_MID, rule.BELL_END)
    res["C_daily_mm"] = daily(t)
    se = {s: _load(wr.DATA / f"{s}.csv.gz") for s in wr.STATIONS}
    tse, _ = table(se, *wr.PRIMARY)
    res["NOT_A_TEST_southeast_this_split"] = arm(tse, rule.WEEK, rule.WKND)
    res["controls"] = c = controls(t, rule.WEEK, rule.WKND, res["primary"]["D"])
    void = []
    if res["stations_with_20_summers"] < rule.MIN_STATIONS: void.append("too few stations")
    if c["excludes_zero"] > 0.15: void.append("interval too narrow on rotated data")
    if c["inject_survives"] < 0.80: void.append("could not see a planted 10 % cycle")
    res["void"] = void
    res["final"] = "VOID" if void else res["primary"]["verdict"]
    (HERE / "results.json").write_text(json.dumps(res, indent=1))
    print(json.dumps(res, indent=1))


if __name__ == "__main__":
    main()
