#!/usr/bin/env python
"""Stage 15 — per-grid-cell binary skill maps of the EHDE forecast, per lead week.

For each model and lead block, pools the 20 held-out years and accumulates the 2x2 contingency
counts PER GRID CELL (using each sample's leakage-safe per-region decision threshold), then
computes POD/POFD/FAR/CSI/EDI/SEDI/HSS/Peirce/freq_bias per cell. Maps let the forecaster see
WHERE the product is skilful at each lead, not just the pooled national number.

Models mapped: the two ensembles (ens_stack, ens_mean) and representative single members
(hgb, mlp). Zoo-member flags are rebuilt from the stored per-fold per-region thresholds at
cascade cutoff 1.0; ensemble flags come from stage 14 (outputs/loyo_pred_v2_ens/ens_pooled.npz).

Outputs:
  outputs/v2_skill_maps.nc                 (model, block, lat, lon) fields for every metric
  results/figures/v2_skillmap_<model>.png  metric(rows) x lead-week(cols) map grid per model
"""
from __future__ import annotations
import os, sys, glob
import numpy as np
import xarray as xr

sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src"))
from ehd.config import load_config, lead_blocks
from ehd import verify as vf

HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.dirname(HERE)
FIGS = os.path.join(ROOT, "results", "figures")
METRICS = ["POD", "POFD", "FAR", "CSI", "EDI", "SEDI", "HSS", "Peirce", "freq_bias"]
PLOT_METRICS = ["POD", "POFD", "FAR", "CSI", "EDI", "SEDI"]
ZOO_MODELS = ["hgb", "mlp"]              # representative sklearn members to map
TORCH_MODELS = ["cnn", "ann_torch"]     # mapped too when their predictions exist
ENS_MODELS = ["ens_stack", "ens_mean"]


def grid_geometry(cfg):
    """Rebuild the land-cell -> (lat,lon) mapping used by features.assemble_year."""
    ref = xr.open_dataarray(os.path.join(cfg["paths"]["outputs"], "forecast_daily_tx",
                                         f"fc_daily_tx_{cfg['forecast']['init_years'][0]}.nc"))
    reg = xr.open_dataset(cfg["paths"]["mask"])["reg"].reindex(
        lat=ref.lat, lon=ref.lon, method="nearest")
    iy, ix = np.where(np.isfinite(reg.values))
    return ref.lat.values, ref.lon.values, iy, ix, reg.values


def counts_per_cell(cell, y, flag, ncell):
    m = np.isfinite(y)
    cell, y, flag = cell[m], y[m].astype(int), flag[m].astype(bool)
    a = np.bincount(cell, weights=(flag & (y == 1)), minlength=ncell)
    b = np.bincount(cell, weights=(flag & (y == 0)), minlength=ncell)
    c = np.bincount(cell, weights=(~flag & (y == 1)), minlength=ncell)
    d = np.bincount(cell, weights=(~flag & (y == 0)), minlength=ncell)
    return a[:ncell], b[:ncell], c[:ncell], d[:ncell]


def metrics_per_cell(a, b, c, d):
    """Vectorised binary scores per cell; NaN where a metric is undefined for that cell."""
    eps = 1e-12
    ev = a + c            # observed events ; nev = b + d observed non-events
    out = {}
    out["POD"] = np.where(ev > 0, a / (ev + eps), np.nan)
    out["POFD"] = np.where((b + d) > 0, b / (b + d + eps), np.nan)
    out["FAR"] = np.where((a + b) > 0, b / (a + b + eps), np.nan)
    out["CSI"] = np.where((a + b + c) > 0, a / (a + b + c + eps), np.nan)
    out["Peirce"] = out["POD"] - out["POFD"]
    out["freq_bias"] = np.where(ev > 0, (a + b) / (ev + eps), np.nan)
    n = a + b + c + d
    hr = (a + d) / (n + eps)
    he = ((a + b) * (a + c) + (c + d) * (b + d)) / (n * n + eps)
    out["HSS"] = np.where(n > 0, (hr - he) / (1 - he + eps), np.nan)
    F = b / (b + d + eps); H = a / (ev + eps)
    with np.errstate(divide="ignore", invalid="ignore"):
        edi = (np.log(F + eps) - np.log(H + eps)) / (np.log(F + eps) + np.log(H + eps) + eps)
        sedi = ((np.log(F + eps) - np.log(H + eps) - np.log(1 - F + eps) + np.log(1 - H + eps))
                / (np.log(F + eps) + np.log(H + eps) + np.log(1 - F + eps)
                   + np.log(1 - H + eps) + eps))
    ok = (ev > 0) & ((b + d) > 0)
    out["EDI"] = np.where(ok, edi, np.nan)
    out["SEDI"] = np.where(ok, sedi, np.nan)
    return out


def main():
    cfg = load_config()
    blocks = lead_blocks(cfg)
    block_names = list(blocks.keys())
    block_ids = {b: k + 1 for k, b in enumerate(block_names)}
    lat, lon, iy, ix, regvals = grid_geometry(cfg)
    ncell = iy.size
    print(f"[maps] grid {lat.size}x{lon.size}, land cells={ncell}")

    files = sorted(glob.glob(os.path.join(ROOT, "outputs", "loyo_pred_v2", "pred_v2_*.npz")))
    cnn_files = {int(os.path.basename(p).split("_")[-1][:-4]): p
                 for p in glob.glob(os.path.join(ROOT, "outputs", "loyo_pred_v2_cnn",
                                                 "pred_v2cnn_*.npz"))}
    single_models = list(ZOO_MODELS) + (list(TORCH_MODELS) if cnn_files else [])
    models = ENS_MODELS + single_models
    print(f"[maps] single members: {single_models}; ensembles: {ENS_MODELS}")

    # accumulate contingency counts per (model, block, cell)
    A = {(m, b): np.zeros(ncell) for m in models for b in block_names}
    B = {k: np.zeros(ncell) for k in A}; C = {k: np.zeros(ncell) for k in A}
    D = {k: np.zeros(ncell) for k in A}

    # --- single members from per-fold predictions + stored per-region thresholds (cutoff 1.0) ---
    for f in files:
        d = np.load(f, allow_pickle=True)
        yr = int(os.path.basename(f).split("_")[-1][:-4])
        dce = np.load(cnn_files[yr], allow_pickle=True) if yr in cnn_files else None
        bid = d["block_id"].astype(int); reg = d["region"].astype(int)
        cell = d["cell"].astype(int); y = d["y"].astype(float)
        for fam in single_models:
            src = d if ("p_" + fam + "_c1") in d.files else dce
            tag = f"{fam}_c1"; p = src["p_" + tag].astype(float); dt = src["decthr"]
            for b in block_names:
                bm = bid == block_ids[b]
                sel = (dt["tag"] == tag) & (dt["block"] == b)
                thr_by_reg = np.full(int(reg.max()) + 1, 0.5)
                for row in dt[sel]:
                    thr_by_reg[int(row["region"])] = float(row["thr"])
                flag = p[bm] >= thr_by_reg[reg[bm]]
                aa, bb, cc, dd = counts_per_cell(cell[bm], y[bm], flag, ncell)
                A[(fam, b)] += aa; B[(fam, b)] += bb; C[(fam, b)] += cc; D[(fam, b)] += dd

    # --- ensembles from stage 14 pooled file (flags already assigned) ---
    ep = np.load(os.path.join(ROOT, "outputs", "loyo_pred_v2_ens", "ens_pooled.npz"),
                 allow_pickle=True)
    ebid = ep["block_id"].astype(int); ecell = ep["cell"].astype(int); ey = ep["y"].astype(float)
    for model, fkey in (("ens_stack", "flag_ensstack"), ("ens_mean", "flag_ensmean")):
        flag_all = ep[fkey].astype(bool)
        for b in block_names:
            bm = ebid == block_ids[b]
            aa, bb, cc, dd = counts_per_cell(ecell[bm], ey[bm], flag_all[bm], ncell)
            A[(model, b)] += aa; B[(model, b)] += bb; C[(model, b)] += cc; D[(model, b)] += dd

    # --- build gridded fields ---
    def to_grid(vec):
        g = np.full((lat.size, lon.size), np.nan)
        g[iy, ix] = vec
        return g

    data_vars = {}
    for metric in METRICS:
        arr = np.full((len(models), len(block_names), lat.size, lon.size), np.nan, np.float32)
        for mi, m in enumerate(models):
            for bi, b in enumerate(block_names):
                sc = metrics_per_cell(A[(m, b)], B[(m, b)], C[(m, b)], D[(m, b)])
                arr[mi, bi] = to_grid(sc[metric])
        data_vars[metric] = (("model", "block", "lat", "lon"), arr)
    ncell_field = np.full((len(models), len(block_names), lat.size, lon.size), np.nan, np.float32)
    for mi, m in enumerate(models):
        for bi, b in enumerate(block_names):
            ncell_field[mi, bi] = to_grid(A[(m, b)] + B[(m, b)] + C[(m, b)] + D[(m, b)])
    data_vars["n_samples"] = (("model", "block", "lat", "lon"), ncell_field)

    ds = xr.Dataset(data_vars, coords=dict(model=models, block=block_names, lat=lat, lon=lon))
    ds.attrs["description"] = ("Per-cell binary skill of the v2 EHDE forecast, pooled over 20 "
                               "LOYO held-out years; leakage-safe per-region decision thresholds.")
    outnc = os.path.join(cfg["paths"]["outputs"], "v2_skill_maps.nc")
    ds.to_netcdf(outnc)
    print(f"[maps] wrote {outnc}")

    _figures(ds, models, block_names)


def _figures(ds, models, block_names):
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    from matplotlib.colors import BoundaryNorm, LinearSegmentedColormap, Normalize

    levels = [round(0.1 * i, 1) for i in range(11)]         # 0.0..1.0, tick spacing 0.1
    # POD/POFD/FAR: smooth (continuous) dark-green(0) -> yellow(0.5) -> dark-red(1) ramp.
    div_cmap = plt.get_cmap("RdYlGn_r")
    div_norm = Normalize(vmin=0.0, vmax=1.0)
    # CSI/EDI/SEDI: yellow (0) -> red (1); negatives (EDI/SEDI can dip <0) clamp to yellow.
    seq_cmap = LinearSegmentedColormap.from_list("yellow_red", ["#FFFF00", "#D00000"])
    seq_cmap.set_under("#FFFF00")
    seq_norm = BoundaryNorm(levels, ncolors=seq_cmap.N)
    styles = dict(POD=(div_cmap, div_norm), POFD=(div_cmap, div_norm), FAR=(div_cmap, div_norm),
                  CSI=(seq_cmap, seq_norm), EDI=(seq_cmap, seq_norm), SEDI=(seq_cmap, seq_norm))
    for m in models:
        nr, nc = len(PLOT_METRICS), len(block_names)
        fig, axes = plt.subplots(nr, nc, figsize=(2.05 * nc, 1.95 * nr),
                                 subplot_kw={}, squeeze=False)
        for ri, metric in enumerate(PLOT_METRICS):
            cmap, norm = styles[metric]
            for ci, b in enumerate(block_names):
                ax = axes[ri][ci]
                g = ds[metric].sel(model=m, block=b).values
                im = ax.pcolormesh(ds.lon, ds.lat, g, cmap=cmap, norm=norm, shading="auto")
                ax.set_xticks([]); ax.set_yticks([])
                if ri == 0:
                    ax.set_title(b, fontsize=9)
                if ci == 0:
                    ax.set_ylabel(metric, fontsize=9)
            cax = fig.add_axes([0.915, axes[ri][-1].get_position().y0,
                                0.012, axes[ri][-1].get_position().height])
            fig.colorbar(im, cax=cax, ticks=levels)
        fig.suptitle(f"v2 EHDE per-cell binary skill — {m}", fontsize=12)
        fig.subplots_adjust(left=0.05, right=0.9, top=0.93, bottom=0.03, wspace=0.05, hspace=0.12)
        out = os.path.join(FIGS, f"v2_skillmap_{m}.png")
        fig.savefig(out, dpi=125); plt.close(fig)
        print(f"[maps] wrote {out}")


if __name__ == "__main__":
    main()