#!/usr/bin/env python
"""Stage 22 — item 4: grid-point probability P(yEHDE | g, Wk) from the member vote fraction.

P(yEHDE | g, Wk) = weighted fraction of members voting yEHDE at cell g in week k, using the
registry-v3 pool and the FAMILY-BALANCED weighting required by <member_definition>: configs within
a family are strongly correlated, so a naive member count lets the biggest family dominate. We give
each ACTIVE family (per-week mask from registry v2) equal total weight, split equally across its
active members:  P = mean_family( mean vote within family ).  We report the nominal member count and
the effective ensemble size  N_eff = (Sum w)^2 / Sum w^2  alongside.

The raw fraction is only a starting point: it is isotonic-recalibrated INSIDE the LOYO folds (fit on
the other 19 years, applied to the held-out year). Both raw and recalibrated reliability are reported
so the effect of recalibration is visible. The per-cell recalibrated mean probability and its Brier
skill vs climatology are saved as CF-compliant NetCDF alongside the operational binary field.

Outputs:
  results/tables/v3_grid_probability_scores.csv   per block: Brier, BSS, N_eff, reliability (raw/recal)
  results/tables/v3_grid_reliability.csv           reliability-curve points (raw & recal) per block
  results/figures/v3_grid_reliability.png          reliability diagrams + sharpness, raw vs recal
  outputs/v3_grid_probability.nc                   prob_ehd(block,lat,lon) recal mean + bss + Neff
"""
from __future__ import annotations
import os, sys, re, logging
import numpy as np
import pandas as pd
import xarray as xr
from sklearn.isotonic import IsotonicRegression

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

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger("stage22")


def family_of(mid):
    for suf in ("_tele_c1", "_bal_c1"):
        if mid.endswith(suf):
            return mid[:-len(suf)]
    return re.sub(r"_c[0-9.]+$", "", mid)


def family_balanced_weights(members, active):
    """Weight vector (sums to 1): equal weight per ACTIVE family, split across its active members."""
    fams = np.array([family_of(m) for m in members])
    w = np.zeros(len(members))
    act_fams = sorted({f for f, a in zip(fams, active) if a})
    for f in act_fams:
        idx = [i for i in range(len(members)) if fams[i] == f and active[i]]
        for i in idx:
            w[i] = 1.0 / len(act_fams) / len(idx)
    return w, act_fams


def main():
    cfg = load_config()
    blocks = lead_blocks(cfg)
    tdir = os.path.join(cfg["paths"]["results"], "tables")
    fdir = os.path.join(cfg["paths"]["results"], "figures")
    log.info("loading v3 member pool ...")
    data, members = mv.load_pool(cfg, blocks, include_v3=True)
    mask = mv.active_mask(cfg, members, blocks)
    years = sorted(data)

    # grid geometry for the NetCDF
    ref = xr.open_dataarray(os.path.join(cfg["paths"]["outputs"], "forecast_daily_tx",
                                         f"fc_daily_tx_{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))
    lat = ref.lat.values; lon = ref.lon.values
    ncell = iy.size

    score_rows, rel_rows = [], []
    grid_prob = np.full((len(blocks), lat.size, lon.size), np.nan, np.float32)
    grid_bss = np.full((len(blocks), lat.size, lon.size), np.nan, np.float32)
    neff_by_block = {}

    for bi, b in enumerate(blocks):
        w, act_fams = family_balanced_weights(members, mask[b])
        n_eff = 1.0 / np.sum(w ** 2)
        n_active = int(np.sum(mask[b]))
        neff_by_block[b] = (n_active, len(act_fams), n_eff)

        # pooled OOS: raw family-balanced fraction, in-CV isotonic recal, base-rate clim
        yA, pA_raw, pA_rec, cellA = [], [], [], []
        for y in years:
            r = data[y][b]
            fr = r["V"].astype(float) @ w                      # weighted vote fraction
            tr = [t for t in years if t != y]
            ftr = np.concatenate([data[t][b]["V"].astype(float) @ w for t in tr])
            ytr = np.concatenate([data[t][b]["y"] for t in tr])
            m = np.isfinite(ftr) & np.isfinite(ytr)
            iso = IsotonicRegression(out_of_bounds="clip").fit(ftr[m], ytr[m].astype(int))
            yA.append(r["y"]); pA_raw.append(fr); pA_rec.append(iso.transform(fr))
            cellA.append(r["cell"])
        y = np.concatenate(yA); p_raw = np.concatenate(pA_raw); p_rec = np.concatenate(pA_rec)
        cell = np.concatenate(cellA)
        clim = np.full_like(y, np.nanmean(y), dtype=float)

        bss_raw = vf.brier_skill_score(y, p_raw, clim)
        bss_rec = vf.brier_skill_score(y, p_rec, clim)
        score_rows.append(dict(block=b, n_active_members=n_active, n_active_families=len(act_fams),
                               N_eff=round(n_eff, 2), base_rate=round(float(np.nanmean(y)), 4),
                               brier_raw=round(vf.brier(y, p_raw), 4),
                               brier_recal=round(vf.brier(y, p_rec), 4),
                               bss_raw=round(bss_raw, 4), bss_recal=round(bss_rec, 4),
                               auc=round(vf.auc(y, p_rec), 4),
                               sharpness_recal=round(float(np.nanstd(p_rec)), 4)))
        for tag, p in (("raw", p_raw), ("recal", p_rec)):
            conf, obs, cnt = vf.reliability_curve(y, p, nbins=10)
            for k in range(len(conf)):
                rel_rows.append(dict(block=b, kind=tag, bin=k, conf=conf[k], obs=obs[k], count=int(cnt[k])))

        # per-cell mean recal prob + BSS (climatological map)
        for gi in range(ncell):
            m = cell == gi
            if m.sum() < 20:
                continue
            grid_prob[bi, iy[gi], ix[gi]] = np.nanmean(p_rec[m])
            cl = np.full(int(m.sum()), np.nanmean(y[m]))
            grid_bss[bi, iy[gi], ix[gi]] = vf.brier_skill_score(y[m], p_rec[m], cl)
        log.info("%s: %d active members (%d families) N_eff=%.1f BSS_recal=%.3f",
                 b, n_active, len(act_fams), n_eff, bss_rec)

    pd.DataFrame(score_rows).to_csv(os.path.join(tdir, "v3_grid_probability_scores.csv"), index=False)
    pd.DataFrame(rel_rows).to_csv(os.path.join(tdir, "v3_grid_reliability.csv"), index=False)

    # NetCDF product
    ds = xr.Dataset(
        {"prob_ehd": (("block", "lat", "lon"), grid_prob,
                      {"long_name": "family-balanced recalibrated P(>=1 EHD in block), LOYO-mean",
                       "units": "1"}),
         "bss_vs_clim": (("block", "lat", "lon"), grid_bss,
                         {"long_name": "Brier skill score vs per-cell climatology", "units": "1"})},
        coords={"block": np.arange(1, len(blocks) + 1),
                "lat": lat.astype("float32"), "lon": lon.astype("float32")},
        attrs={"Conventions": "CF-1.8",
               "title": "Stage-2 grid vote-fraction EHD probability (registry v3, family-balanced)",
               "comment": ("P(yEHDE|g,Wk) = family-balanced member vote fraction, isotonic-recalibrated "
                           "in-CV; block index 1..6 = " + " ".join(blocks)),
               "n_eff_by_block": " ".join(f"{b}:{neff_by_block[b][2]:.1f}" for b in blocks)})
    ds.to_netcdf(os.path.join(cfg["paths"]["outputs"], "v3_grid_probability.nc"))

    # reliability figure
    import matplotlib; matplotlib.use("Agg"); import matplotlib.pyplot as plt
    rel = pd.DataFrame(rel_rows)
    fig, axes = plt.subplots(2, 3, figsize=(13, 8))
    for ax, b in zip(axes.ravel(), blocks):
        for tag, col in (("raw", "0.6"), ("recal", "tab:red")):
            sub = rel[(rel.block == b) & (rel.kind == tag)]
            ax.plot(sub.conf, sub.obs, "o-", color=col, label=tag, ms=4)
        ax.plot([0, 1], [0, 1], "k--", lw=.8)
        na, nf, ne = neff_by_block[b]
        ax.set_title(f"{b}  (N={na}, N_eff={ne:.1f})", fontsize=10)
        ax.set_xlabel("forecast prob"); ax.set_ylabel("observed freq"); ax.legend(fontsize=8)
    fig.suptitle("Stage-22 grid vote-fraction reliability — raw vs in-CV isotonic recalibration")
    fig.tight_layout(); fig.savefig(os.path.join(fdir, "v3_grid_reliability.png"), dpi=130)

    print(pd.DataFrame(score_rows).to_string(index=False))


if __name__ == "__main__":
    main()