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

sys.path.insert(0, str(Path(__file__).parent))
import rule  # noqa: E402

HERE = Path(__file__).parent


def table(all_data, y0, y1):
    """{year: [[sum_tenths, n, wet] for weekday 0..6]}, plus counts."""
    by_year, dropped = rule.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, shift=0, wet=False, boost=1.0):
    """(mid_sum, mid_n, end_sum, end_n) for one year, labels rotated by shift."""
    m = nm = e = ne = 0.0
    k = 2 if wet else 0
    for wd in range(7):
        lab = (wd + shift) % 7
        if lab in rule.MID:
            m += row[wd][k] * boost; nm += row[wd][1]
        elif lab in rule.END:
            e += row[wd][k]; ne += row[wd][1]
    return (m, nm, e, ne)


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


def weekdays(t):
    names = "Mon Tue Wed Thu Fri Sat Sun".split()
    ys = sorted(t)
    rng = random.Random(20261008)
    def mean(sel, wd):
        s = sum(t[y][wd][0] for y in sel); n = sum(t[y][wd][1] for y in sel)
        return s / n / 10
    out = {}
    for wd in range(7):
        vals = sorted(mean([rng.choice(ys) for _ in ys], wd) for _ in range(2000))
        out[names[wd]] = {"mm": mean(ys, wd), "lo": vals[50], "hi": vals[1949]}
    return out


def controls(t, real_d):
    rng = random.Random(7)
    ys = sorted(t)
    excl = surv = ge = 0
    inj = inj_surv = 0
    for i in range(2000):
        sh = {y: rng.randrange(7) for y in ys}
        yp = {y: parts(t[y], shift=sh[y]) for y in ys}
        d = rule.D(list(yp.values()))
        lo, hi = rule.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], shift=sh[y], boost=1.10) for y in ys}
            db = rule.D(list(yb.values()))
            lob, hib = rule.boot(yb, 500, 5000 + i)
            inj += 1
            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_draws": inj,
            "inject_survives": inj_surv / inj}


def main():
    all_data = {}
    missing = []
    for sid in rule.STATIONS:
        f = rule.DATA / f"{sid}.csv.gz"
        if not f.exists():
            missing.append(sid); continue
        all_data[sid] = rule.load(sid)
    res = {"stations_loaded": len(all_data), "missing_files": missing}
    t, dropped = table(all_data, *rule.PRIMARY)
    good = 0
    for sid, rec in all_data.items():
        n = sum(1 for y in range(rule.PRIMARY[0], rule.PRIMARY[1] + 1)
                if sum(1 for d in rec if d.year == y) >= 85)
        good += n >= 20
    res["stations_with_20_summers"] = good
    res["primary_dropped_station_summers"] = dropped
    res["station_days_primary"] = sum(c[1] for r in t.values() for c in r)
    res["primary"] = arm(t)
    res["A_wetday"] = arm(t, wet=True)
    for k, (y0, y1) in rule.ARMS.items():
        tk, dk = table(all_data, y0, y1)
        res[k] = arm(tk); res[k]["dropped"] = dk
    res["D_weekdays_mm"] = weekdays(t)
    res["controls"] = controls(t, res["primary"]["D"])
    c = res["controls"]
    void = []
    if good < 40: void.append("fewer than 40 stations with 20 summers")
    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()