#!/usr/bin/env python3
"""Score EHDE as a regional event rather than a grid-point event.

A single grid point turning EHD=1 three to six weeks ahead is close to
unpredictable.  The framing here, proposed by the user, keeps the forecast at
grid-point resolution but moves the verification up to the sub-region: for each
of the seven zones on each valid day,

    P_i = m_i / n_i        m_i = cells with EHD = 1,  n_i = cells in region i

and the region counts as an EHDE day when P_i exceeds a coverage threshold.

Two ways of turning the forecast into the same yes/no answer are reported:

  literal   count the cells the point model puts above its own threshold and
            apply the same coverage rule -- exactly the rule as proposed
  tuned     use the mean predicted probability over the region as a continuous
            regional score, with its decision threshold read off the
            out-of-fold field inside the training years

The tuned route needs a spatially complete out-of-fold field, which the thinned
`oof_*.npz` cannot give (prep drops 85% of non-events, so a coverage fraction
computed from it would be wrong).  It therefore reads `oofpred_*.npz`, written
by `train_national.py predict_oof`, where every training row is scored by the
fold model that held its year out.  No threshold ever sees a validation year.
"""

from __future__ import annotations

import argparse
import glob
import json
import os
import sys

import numpy as np

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import ehd_common as C
import ehd_metrics as M
from train_national import RES_DIR, shard_dir

NREG = 7
KEYS = ("score", "evt", "week", "lt", "init", "region", "capable")


def load(tag, pattern):
    files = sorted(glob.glob(os.path.join(shard_dir(tag), pattern)))
    if not files:
        raise SystemExit("no %s for tag '%s'" % (pattern, tag))
    parts = [np.load(f) for f in files]
    miss = [k for k in KEYS if k not in parts[0].files]
    if miss:
        raise SystemExit("%s lacks %s" % (pattern, ", ".join(miss)))
    d = {k: np.concatenate([p[k] for p in parts]) for k in KEYS}
    ok = np.isfinite(d["score"]) & np.isfinite(d["evt"])
    return {k: v[ok] for k, v in d.items()}


def regionalise(d, npts, point_th=None):
    """Collapse (init, lead, point) rows into one row per (init, lead, region)."""
    key = (d["init"].astype("i8") * C.NLT + d["lt"].astype("i8")) * 8 + d["region"]
    o = np.argsort(key, kind="stable")
    k = key[o]
    cut = np.flatnonzero(np.diff(k))
    st = np.r_[0, cut + 1]
    en = np.r_[cut + 1, k.size]
    reg = (k[st] % 8).astype(int)
    lead = ((k[st] // 8) % C.NLT).astype(int)
    n = npts[reg - 1].astype("f8")

    out = {"region": reg, "lead": lead,
           "week": np.searchsorted(np.array([11, 18, 25, 32, 39]), lead, side="right"),
           "P_obs": np.add.reduceat((d["evt"][o] > 0.5).astype("f8"), st) / n,
           "Pbar": np.add.reduceat(d["score"][o].astype("f8"), st) / (en - st),
           "P_cap": np.add.reduceat((d["capable"][o] > 0.5).astype("f8"), st) / n}
    if point_th is not None:
        hit = np.zeros(d["score"].size, bool)
        for w, t in point_th.items():
            m = d["week"] == int(w)
            hit[m] = d["score"][m] >= t
        out["P_fc"] = np.add.reduceat(hit[o].astype("f8"), st) / n
    return out


MIN_EVT_REGION = 40


def region_week_cut(score, obs, week, region, pooled, rule):
    """One threshold per (region, week), falling back to the pooled week cut.

    Section Y7 measured this at grid-point level: a single national cut
    over-forecasts by a factor 2.2 in S7, the region the model reads best.  The
    same imbalance carries over to the regional score, so give each region its
    own cut wherever it has enough events to place one.
    """
    th = {}
    for r in range(1, NREG + 1):
        for w in range(len(C.WEEKS)):
            m = (region == r) & (week == w)
            if m.sum() and obs[m].sum() >= MIN_EVT_REGION:
                th[(r, w)] = weekly_cut(score[m], obs[m], np.zeros(int(m.sum()), int),
                                        rule)[0]
            else:
                th[(r, w)] = pooled[w]
    return th


def apply_cut(score, week, region, th):
    p = np.zeros(score.size, bool)
    if isinstance(next(iter(th)), tuple):
        for (r, w), t in th.items():
            m = (region == r) & (week == w)
            p[m] = score[m] >= t
    else:
        for w, t in th.items():
            m = week == w
            p[m] = score[m] >= t
    return p


def by_region_table(pred, obs, region, score, title):
    """Same scores, but split by sub-region -- the national row hides a lot."""
    rows = {}
    print("\n  %s" % title)
    print("    %-18s %9s %7s %7s %7s %7s %7s %7s %7s"
          % ("vùng", "base_rate", "n_evt", "POD", "POFD", "FAR", "FBI", "CSI", "AUC"))
    for r in range(1, NREG + 1):
        m = region == r
        if m.sum() == 0 or obs[m].sum() == 0:
            continue
        s = M.score_binary(pred[m], obs[m])
        s["AUC"] = M.auc_score(score[m], obs[m])
        s["n_events"] = int(obs[m].sum())
        rows[C.REGION_NAMES[r]] = s
        print("    %-18s %9.4f %7d %7.3f %7.3f %7.3f %7.3f %7.3f %7.3f"
              % (C.REGION_NAMES[r], s["base_rate"], s["n_events"], s["POD"],
                 s["POFD"], s["FAR"], s["FBI"], s["CSI"], s["AUC"]))
    return rows


def weekly_cut(score, obs, week, rule):
    """One threshold per forecast week, read off the training-year field."""
    th = {}
    for w in range(len(C.WEEKS)):
        m = week == w
        if m.sum() == 0 or obs[m].sum() == 0:
            th[w] = np.inf
            continue
        s, o = score[m], obs[m]
        if rule == "fbi1":
            th[w] = float(np.quantile(s, 1.0 - o.mean()))
        else:
            grid = np.unique(np.quantile(s, np.linspace(0.50, 0.9995, 400)))
            best, bt = -np.inf, float(np.median(s))
            for t in grid:
                f = s >= t
                v = M.scores(float((f & o).sum()), float((f & ~o).sum()),
                             float((~f & o).sum()), float((~f & ~o).sum()))["CSI"]
                if np.isfinite(v) and v > best:
                    best, bt = v, float(t)
            th[w] = bt
    return th


def score_table(pred, obs, week, title, extra=None):
    rows = []
    print("\n  %s" % title)
    print("        %9s %7s %7s %7s %7s %7s %7s %7s"
          % ("base_rate", "n_evt", "POD", "POFD", "FAR", "FBI", "CSI", "AUC"))
    for w in list(range(len(C.WEEKS))) + ["ALL"]:
        m = (week == w) if w != "ALL" else np.ones(obs.size, bool)
        if m.sum() == 0 or obs[m].sum() == 0:
            continue
        f, o = pred[m], obs[m]
        s = M.scores(float((f & o).sum()), float((f & ~o).sum()),
                     float((~f & o).sum()), float((~f & ~o).sum()))
        s["AUC"] = M.auc_score(extra[m], o) if extra is not None else np.nan
        s["base_rate"] = float(o.mean())
        s["n_events"] = int(o.sum())
        name = C.WEEK_NAMES[w] if w != "ALL" else "ALL"
        rows.append((name, s))
        print("    %-4s %9.4f %7d %7.3f %7.3f %7.3f %7.3f %7.3f %7.3f"
              % (name, s["base_rate"], s["n_events"], s["POD"], s["POFD"],
                 s["FAR"], s["FBI"], s["CSI"], s["AUC"]))
    return {n: s for n, s in rows}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--tag", default="_ens")
    ap.add_argument("--name", default="ensemble_F3_ens")
    ap.add_argument("--covers", type=float, nargs="+", default=[0.05, 0.10, 0.20])
    ap.add_argument("--cap-min", type=float, default=0.20,
                    help="a region-day counts as in season when this fraction of its "
                         "cells are heat-capable")
    args = ap.parse_args()

    grid = C.Grid()
    npts = np.array([(grid.reg == r).sum() for r in range(1, NREG + 1)])
    with open(os.path.join(RES_DIR, "national_%s.json" % args.name)) as fh:
        point_th = json.load(fh)["thresholds_fbi1"]

    va = regionalise(load(args.tag, "pred_*.npz"), npts, point_th)
    tr = regionalise(load(args.tag, "oofpred_*.npz"), npts)
    print("[regional] %d region-days kiểm định, %d region-days huấn luyện (out-of-fold)"
          % (va["region"].size, tr["region"].size), flush=True)
    print("[regional] ô lưới mỗi vùng: %s"
          % dict(zip(["S%d" % r for r in range(1, 8)], npts.tolist())), flush=True)

    out = {"cap_min": args.cap_min, "results": {}}
    for cov in args.covers:
        ov, ot = va["P_obs"] > cov, tr["P_obs"] > cov
        for lbl, sv, stt in (("mọi ngày", np.ones(ov.size, bool), np.ones(ot.size, bool)),
                             ("mùa nóng (≥%.0f%% ô capable)" % (100 * args.cap_min),
                              va["P_cap"] > args.cap_min, tr["P_cap"] > args.cap_min)):
            head = "phủ > %.0f%% · %s" % (100 * cov, lbl)
            print("\n" + "=" * 78 + "\n  %s" % head)
            key = "cover%02d_%s" % (int(100 * cov), "season" if "mùa" in lbl else "all")

            if cov == 0.20 and sv.all():
                out["results"][key + "_literal"] = score_table(
                    va["P_fc"][sv] > cov, ov[sv], va["week"][sv],
                    "quy tắc nguyên văn: đếm ô vượt ngưỡng điểm, rồi phủ > 20%",
                    va["P_fc"][sv])

            w_, s_, r_ = va["week"][sv], va["Pbar"][sv], va["region"][sv]
            for rule in ("fbi1", "csi"):
                th = weekly_cut(tr["Pbar"][stt], ot[stt], tr["week"][stt], rule)
                out["results"][key + "_" + rule] = score_table(
                    apply_cut(s_, w_, r_, th), ov[sv], w_,
                    "xác suất trung bình vùng, ngưỡng dò trên OOF (%s)"
                    % ("FBI≈1" if rule == "fbi1" else "tối đa CSI"), s_)
                if rule != "fbi1":
                    continue
                thr = region_week_cut(tr["Pbar"][stt], ot[stt], tr["week"][stt],
                                      tr["region"][stt], th, rule)
                pv = apply_cut(s_, w_, r_, thr)
                out["results"][key + "_fbi1_regweek"] = score_table(
                    pv, ov[sv], w_,
                    "xác suất trung bình vùng, ngưỡng riêng theo (VÙNG, tuần), FBI≈1", s_)
                late = w_ >= 3
                if late.sum() and ov[sv][late].sum():
                    out["results"][key + "_fbi1_regweek_byregion"] = by_region_table(
                        pv[late], ov[sv][late], r_[late], s_[late],
                        "ngưỡng riêng theo (vùng, tuần) — theo vùng, trung bình W4–W6")
                    out["results"][key + "_fbi1_byregion"] = by_region_table(
                        apply_cut(s_, w_, r_, th)[late], ov[sv][late], r_[late], s_[late],
                        "ngưỡng chung theo tuần — theo vùng, trung bình W4–W6")

    path = os.path.join(RES_DIR, "regional_%s.json" % args.name)
    with open(path, "w") as fh:
        json.dump(out, fh, indent=1, default=float)
    print("\n[done] -> %s" % path, flush=True)


if __name__ == "__main__":
    main()