# SPDX-License-Identifier: MIT
"""YS-EN-0006 · score saved runs: MAE by day type, paired bootstrap against a reference.

  python analyse.py results/vtt_main.npz [reference_config]
  python analyse.py stats      # descriptive statistics of the VTT rooms quoted in the note

Per (series, day) MAE over the 24 forecast hours; neural configs average the per-day MAE
over seeds (seed spread reported separately).  Confidence intervals resample whole
calendar dates (all series on a date together), because rooms share holidays and weather.
"""
from __future__ import annotations

import json
import sys

import numpy as np

B = 2000


def load(path):
    z = np.load(path, allow_pickle=False)
    truth = z["truth"]
    lab = {k[5:]: z[k] for k in z.files if k.startswith("lab__")}
    preds = {k[6:]: z[k] for k in z.files if k.startswith("pred__")}
    meta = json.loads(str(z["meta"]))
    return truth, preds, lab, meta


def day_mae(pred, truth):
    """[seeds, N, 24] → per-day MAE averaged over seeds [N], and per-seed overall MAE."""
    e = np.abs(pred - truth[None]).mean(-1)
    return e.mean(0), e.mean(1)


def day_types(lab):
    wd, hol, prev = lab["workday"], lab["holiday"], lab["prev_workday"]
    return {
        "all": np.ones_like(wd, bool),
        "workday": wd == 1,
        "first workday back": (wd == 1) & (prev == 0),
        "weekend": (wd == 0) & (hol == 0),
        "holiday on a weekday": (hol == 1) & (lab["dow"] < 5),
    }


def boot_diff(a, b, dates, rng):
    """Mean of (a - b) and 95% CI, resampling dates."""
    u, inv = np.unique(dates, return_inverse=True)
    sa = np.bincount(inv, a, len(u))
    sb = np.bincount(inv, b, len(u))
    n = np.bincount(inv, minlength=len(u)).astype(float)
    idx = rng.integers(0, len(u), (B, len(u)))
    d = (sa[idx].sum(1) - sb[idx].sum(1)) / n[idx].sum(1)
    return float((a - b).mean()), float(np.percentile(d, 2.5)), float(np.percentile(d, 97.5))


def summary(path, ref=None):
    truth, preds, lab, meta = load(path)
    types = day_types(lab)
    rng = np.random.default_rng(0)
    per_day = {k: day_mae(v, truth) for k, v in preds.items()}
    out = {"n_days": {k: int(m.sum()) for k, m in types.items()}, "configs": {}}
    for name, (dm, seeds) in per_day.items():
        row = {t: float(dm[m].mean()) if m.any() else None for t, m in types.items()}
        row["seed_sd"] = float(seeds.std()) if len(seeds) > 1 else 0.0
        info = meta.get(name, [{}])[0]
        if "params" in info:
            row["params"] = info["params"]
        if ref and name != ref:
            mean, lo, hi = boot_diff(dm, per_day[ref][0], lab["date"], rng)
            row["diff_vs_ref"] = [mean, lo, hi]
        out["configs"][name] = row
    return out


def show(out, ref=None):
    types = list(out["n_days"])
    print("days: " + ", ".join(f"{t} {n}" for t, n in out["n_days"].items()))
    print(f"{'config':22s}" + "".join(f"{t[:14]:>15s}" for t in types) + "    seed sd"
          + (f"   Δ vs {ref} [95% CI]" if ref else ""))
    for name, r in out["configs"].items():
        cells = "".join(f"{r[t]:15.2f}" if r[t] is not None else f"{'-':>15s}" for t in types)
        d = r.get("diff_vs_ref")
        tail = f"   {d[0]:+7.2f} [{d[1]:+.2f}, {d[2]:+.2f}]" if d else ""
        print(f"{name:22s}{cells}{r['seed_sd']:11.3f}{tail}")


def stats():
    """Every descriptive number about the VTT data quoted in the note, with its definition."""
    import exp_vtt as V
    S = V.series()
    truth, preds, lab, meta = load(V.OUT / "vtt_main.npz")
    keys = np.load(V.OUT / "vtt_main.npz")["keys"]
    T = day_types(lab)
    out = {}
    first = []
    for s in S:
        ok = np.where(np.isfinite(s["y"]))[0]
        first.append(str(s["start"] + __import__("datetime").timedelta(int(ok[0]) // 24)))
    out["first_reading_per_room"] = dict(zip([s["name"] for s in S], first))
    lo, hi = V.TR
    night, midday = [], []
    for s in S:
        y, wd = s["y"][lo:hi], s["workday"][lo:hi]
        h = np.arange(lo, hi) % 24
        night.append(float(np.nanmedian(y[(h >= 1) & (h <= 4)])))
        midday.append(float(np.nanmedian(y[(wd == 1) & (h >= 10) & (h <= 15)])))
    out["training: per-room median CO2 01:00-05:00, ppm (min, max)"] = [min(night), max(night)]
    out["training: per-room median CO2 10:00-16:00 on workdays, ppm (min, max)"] = [min(midday), max(midday)]
    rng_ = truth.max(1) - truth.min(1)
    out["test workdays: median daily range (max - min hour), ppm"] = float(np.median(rng_[T["workday"]]))
    # flat forecast: the median of the history's 01:00-05:00 readings, held for all 24 hours
    flat = np.array([np.full(24, np.nanmedian(S[i]["y"][t - 168:t][(np.arange(t - 168, t) % 24 >= 1)
                                                                    & (np.arange(t - 168, t) % 24 <= 4)]))
                     for i, t in keys])
    e = np.abs(flat - truth).mean(1)
    out["flat night-level forecast MAE, ppm"] = {t: float(e[m].mean()) for t, m in T.items()}
    per_room = np.bincount(keys[:, 0])
    out["test room-days per room (min, max)"] = [int(per_room.min()), int(per_room.max())]
    return out


if __name__ == "__main__":
    if sys.argv[1] == "stats":
        o = stats()
        print(json.dumps(o, indent=1))
        with open("results/vtt_stats.json", "w") as f:
            json.dump(o, f, indent=1)
        sys.exit()
    path = sys.argv[1]
    ref = sys.argv[2] if len(sys.argv) > 2 else None
    o = summary(path, ref)
    show(o, ref)
    with open(path.replace(".npz", "_summary.json"), "w") as f:
        json.dump(o, f, indent=1)