"""Run the sealed rule on ~/data/ice-saints/daily.parquet, write results.json."""
import datetime as dt, json, sys
from pathlib import Path
import numpy as np, pandas as pd
sys.path.insert(0, str(Path(__file__).parent))
import rule

df = pd.read_parquet(Path.home() / "data/ice-saints/daily.parquet")
N = dt.date(rule.Y1, 12, 31).toordinal() + 40 - rule.ORIGIN
ords = np.array([d.toordinal() for d in df.date.dt.date]) - rule.ORIGIN
keep = (ords >= 0) & (ords < N)
df, ords = df[keep], ords[keep]
arrays, frost = {}, {}
for sid in rule.FRAME:
    m = (df.station == sid).values
    a = np.full(N, np.nan); a[ords[m]] = df.tmk.values[m]
    b = np.full(N, np.nan); b[ords[m]] = df.tnk.values[m]
    arrays[sid], frost[sid] = a, b
A = list(arrays.values())

allv = np.concatenate([a[~np.isnan(a)] for a in A])
sta = {}
for l in open(Path.home() / "data/ice-saints/stations.txt", encoding="latin1").read().split("\n")[2:]:
    p = l.split()
    if len(p) > 5 and p[0] in rule.FRAME:
        sta[p[0]] = dict(elev=float(p[3]), lat=float(p[4]), name=" ".join(p[6:-2]))

yv = rule.year_values(A, 5, 11)
c = float(np.mean(list(yv.values())))
pl = rule.placebos()
pv = {f"{m:02d}-{d:02d}": rule.cbar(A, m, d) for m, d in pl}
pv_a = np.array(list(pv.values()))
boot = rule.bootstrap_se(list(yv.values()))
sd = float(pv_a.std(ddof=1))
se = max(boot, sd)
v, lo, hi = rule.verdict(c, se)

# gates
y29 = rule.year_values(A, 2, 10, y0=1929, y1=1929).get(1929, float("nan"))
def month_mean(mo):
    vals = []
    for a in A:
        for y in range(rule.Y0, rule.Y1 + 1):
            s = dt.date(y, mo, 1).toordinal() - rule.ORIGIN
            vals.append(a[s : s + 28])
    return float(np.nanmean(np.concatenate(vals)))
jj = month_mean(7) - month_mean(1)
planted = {}
for m, d in pl[:: len(pl) // 10][:10]:
    p_c = rule.cbar(A, m, d) + rule.PLANT
    planted[f"{m:02d}-{d:02d}"] = (round(p_c, 3), rule.verdict(p_c, se)[0])
n_surv = sum(1 for x in planted.values() if x[1] == "survived")
gates = {
    "V1": [2.8 * se, 2.8 * se < rule.BAR],
    "V2": [float(y29), bool(y29 > 1.0)],
    "V3": [jj, bool(jj > 12)],
    "V4": [n_surv, n_surv >= 8],
    "V5": [dict(years=len(yv), stations=sum(1 for a in A if (~np.isnan(a)).any()), min=float(allv.min()), max=float(allv.max())),
           bool(len(yv) >= 90 and sum(1 for a in A if (~np.isnan(a)).any()) >= 25 and allv.min() >= -40 and allv.max() <= 45)],
}
void = [k for k, x in gates.items() if not x[1]]
final = "void" if void else v

# printed only
def span(y0, y1, arrs=A):
    vals = list(rule.year_values(arrs, 5, 11, y0=y0, y1=y1).values())
    return dict(c=round(float(np.mean(vals)), 3), se_boot=round(rule.bootstrap_se(vals), 3), years=len(vals))
north = [arrays[s] for s in rule.FRAME if sta[s]["lat"] >= 51]
south = [arrays[s] for s in rule.FRAME if sta[s]["lat"] < 51]
def spanmin(arrs):
    o = []
    for y in range(rule.Y0, rule.Y1 + 1):
        s = dt.date(y, 5, 11).toordinal()
        x = np.array([rule.station_value(a, s) for a in arrs], float)
        if (~np.isnan(x)).sum() >= max(3, len(arrs) // 2):
            o.append(np.nanmean(x))
    return dict(c=round(float(np.mean(o)), 3), years=len(o), stations=len(arrs))
rank = int((pv_a > c).sum()) + 1
sheep = rule.cbar(A, 6, 9)
low = [s for s in rule.FRAME if sta[s]["elev"] <= 400]
fro = {}
for y0, y1 in ((1926, 1955), (1956, 1985), (1986, 2015), (1996, 2025)):
    hit = {k: 0 for k in ("after_may_1", "after_may_16", "after_may_26")}
    tot = 0
    for s in low:
        b = frost[s]
        for y in range(y0, y1 + 1):
            i0 = dt.date(y, 5, 1).toordinal() - rule.ORIGIN
            i1 = dt.date(y, 7, 1).toordinal() - rule.ORIGIN
            x = b[i0:i1]
            if np.isnan(x).sum() > 3:
                continue
            tot += 1
            hit["after_may_1"] += bool(np.nansum(x < 0) > 0)
            hit["after_may_16"] += bool(np.nansum(x[15:] < 0) > 0)
            hit["after_may_26"] += bool(np.nansum(x[25:] < 0) > 0)
    fro[f"{y0}-{y1}"] = dict(station_years=tot, **{k: round(n / tot, 4) for k, n in hit.items()})

# daily line for the chart, May 1 to Jun 5, all counted years, mean T relative to a straight line
out = dict(verdict=final, raw_verdict=v, void=void, C=round(c, 4), lo=round(lo, 4), hi=round(hi, 4), se=round(se, 4),
           se_boot=round(boot, 4), se_placebos=round(sd, 4), years=len(yv), year_sd=round(float(np.std(list(yv.values()), ddof=1)), 3),
           gates=gates, placebos={k: round(x, 3) for k, x in pv.items()}, planted=planted,
           print_only=dict(y1951_2025=span(1951, 2025), y1996_2025=span(1996, 2025), north=spanmin(north), south=spanmin(south),
                           rank_among_placebos=f"{rank} of {len(pv_a)+1}", placebo_max=round(float(pv_a.max()), 3), placebo_min=round(float(pv_a.min()), 3),
                           sheep_cold_jun9=round(sheep, 3), frost_low_stations=fro, low_stations=len(low)),
           year_values={str(k): round(x, 3) for k, x in yv.items()})
json.dump(out, open(Path(__file__).parent / "results.json", "w"), indent=1, default=str)
print({k: out[k] for k in ("verdict", "raw_verdict", "void", "C", "lo", "hi", "se", "se_boot", "se_placebos", "years", "year_sd", "gates", "planted", "print_only")})