"""Turn the raw hourly JSON into one row per (site, day), by rule.py's definitions.

Shared by controls.py and analyse.py so that both score the same days.
"""
import glob
import json
import os

import numpy as np

OUT = os.path.expanduser("~/data/rain-before-seven")
SITES = ["london", "manchester", "edinburgh", "cardiff", "belfast", "plymouth"]
WET_MM = 0.1          # rule.py: an hour is WET at >= 0.1 mm
START, END = "2000-01-01", "2025-12-31"


def _site_hours(site):
    """{'YYYY-MM-DDTHH:00': mm} for one site, over every chunk file."""
    hours = {}
    for path in sorted(glob.glob(os.path.join(OUT, f"{site}_*.json"))):
        with open(path) as f:
            d = json.load(f)
        for t, p in zip(d["hourly"]["time"], d["hourly"]["precipitation"]):
            hours[t] = p
    return hours


def days(wet_mm=WET_MM):
    """One record per (site, day) with the flags rule.py defines.

    A day is dropped if any hour 00:00..10:59 is missing or null.
    Returns a dict of numpy arrays, all the same length.
    """
    site_idx, block, date_s, R, F, Fs, dryh = [], [], [], [], [], [], []
    for si, site in enumerate(SITES):
        hours = _site_hours(site)
        dates = sorted({t[:10] for t in hours})
        for d in dates:
            if not (START <= d <= END):
                continue
            vals = [hours.get(f"{d}T{h:02d}:00") for h in range(15)]
            if any(v is None for v in vals[:11]):
                continue
            # hours 11..14 are descriptive only and may be missing
            dryh.append([np.nan if v is None else float(v < wet_mm)
                         for v in vals])
            vals = vals[:11]
            site_idx.append(si)
            block.append(f"{site}|{d[:7]}")          # (site, calendar month)
            date_s.append(d)
            R.append(any(v >= wet_mm for v in vals[0:7]))     # 00:00..06:59
            F.append(vals[10] < wet_mm)                       # 10:00..10:59 dry
            Fs.append(all(v < wet_mm for v in vals[8:11]))    # 08:00..10:59 dry
    blocks = sorted(set(block))
    bidx = {b: i for i, b in enumerate(blocks)}
    return {
        "site": np.array(site_idx),
        "block": np.array([bidx[b] for b in block]),
        "n_blocks": len(blocks),
        "date": np.array(date_s),
        "R": np.array(R),
        "F": np.array(F),
        "Fstrict": np.array(Fs),
        "dryh": np.array(dryh),          # (days, 15) 1 dry / 0 wet / nan missing
    }


def hourly_dry(d):
    """{hour: (share dry given R, share dry given not R)} for hours 0..14."""
    R = d["R"]
    out = {}
    for h in range(d["dryh"].shape[1]):
        col = d["dryh"][:, h]
        out[h] = (float(np.nanmean(col[R])), float(np.nanmean(col[~R])))
    return out


def delta(R, F):
    """P(F | R) - P(F | not R), in percentage points. nan if either side empty."""
    if R.sum() == 0 or (~R).sum() == 0:
        return float("nan")
    return 100.0 * (F[R].mean() - F[~R].mean())


def delta_strat(block, n_blocks, R, F):
    """Mantel-Haenszel risk difference: delta computed WITHIN each block and
    then averaged, weights n_R*n_notR/n. Unlike the pooled difference it
    cannot be moved by which blocks happen to hold more rainy mornings.

    NOT pre-registered. Added after the null control showed the pooled
    difference carries a between-month composition term; reported as a
    follow-up, never as the sealed headline.
    """
    Fv = F.astype(float)
    n1 = np.bincount(block[R], minlength=n_blocks).astype(float)
    n0 = np.bincount(block[~R], minlength=n_blocks).astype(float)
    s1 = np.bincount(block[R], weights=Fv[R], minlength=n_blocks)
    s0 = np.bincount(block[~R], weights=Fv[~R], minlength=n_blocks)
    ok = (n1 > 0) & (n0 > 0)
    if not ok.any():
        return float("nan")
    w = n1[ok] * n0[ok] / (n1[ok] + n0[ok])
    diff = s1[ok] / n1[ok] - s0[ok] / n0[ok]
    return 100.0 * float((w * diff).sum() / w.sum())


def block_bootstrap(d, R, F, reps=10000, seed=20260919, stat="pooled"):
    """95 % CI for delta, resampling (site, month) blocks with replacement.

    stat="pooled"  -> the sealed statistic, P(F|R) - P(F|not R) over all days
    stat="strat"   -> the within-block (Mantel-Haenszel) version
    """
    rng = np.random.default_rng(seed)
    order = np.argsort(d["block"], kind="stable")
    b_sorted = d["block"][order]
    starts = np.searchsorted(b_sorted, np.arange(d["n_blocks"]), side="left")
    ends = np.searchsorted(b_sorted, np.arange(d["n_blocks"]), side="right")
    members = [order[s:e] for s, e in zip(starts, ends)]
    out = np.empty(reps)
    for i in range(reps):
        pick = rng.integers(0, d["n_blocks"], d["n_blocks"])
        idx = np.concatenate([members[j] for j in pick])
        if stat == "pooled":
            out[i] = delta(R[idx], F[idx])
        else:
            # each drawn block keeps its own identity, re-labelled 0..n-1 so a
            # block drawn twice is counted twice
            rb = np.concatenate([np.full(len(members[j]), k)
                                 for k, j in enumerate(pick)])
            out[i] = delta_strat(rb, len(pick), R[idx], F[idx])
    out = out[~np.isnan(out)]
    return float(np.percentile(out, 2.5)), float(np.percentile(out, 97.5))