#!/usr/bin/env python
"""Stage 23 — item 5: grid-point event STACK vs COUNTING (does stacking beat the vote fraction?).

Base learners = the registry-v3 members. Meta-learners are trained on OUT-OF-FOLD base predictions
(the members' own LOYO out-of-sample probabilities are exactly OOF), fit with an inner LOYO at the
meta level: for held-out year y the meta-model is fit on the other 19 years' (member_probs -> y) and
applied to y. Using in-fold base predictions would be the classic stacking leak; it is avoided here.

Meta-learners compared, all against item-4's COUNTING baseline (family-balanced vote fraction,
in-CV isotonic-recalibrated):
  stack_logit  penalised logistic regression on member probabilities
  stack_hgb    histogram gradient boosting on member probabilities
  bma          weighted (Brier-optimal) model average — the interpretable reference

The honest question: if stacking does not beat counting, recommend the simpler model. Reported per
lead block with BSS vs climatology, AUC, and SEDI/Peirce at the in-CV Peirce-optimal threshold.

Outputs: results/tables/v3_stack_vs_count.csv ; results/figures/v3_stack_vs_count.png
"""
from __future__ import annotations
import os, sys, logging
import numpy as np
import pandas as pd
from sklearn.linear_model import LogisticRegression
from sklearn.preprocessing import StandardScaler
from sklearn.pipeline import make_pipeline
from sklearn.ensemble import HistGradientBoostingClassifier
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("stage23")


def _fit_apply(kind, Xtr, ytr, Xte):
    m = np.isfinite(ytr) & np.all(np.isfinite(Xtr), axis=1)
    if kind == "stack_logit":
        mdl = make_pipeline(StandardScaler(), LogisticRegression(C=1.0, max_iter=1000))
        mdl.fit(Xtr[m], ytr[m].astype(int))
        return mdl.predict_proba(np.nan_to_num(Xte))[:, 1]
    if kind == "stack_hgb":
        mdl = HistGradientBoostingClassifier(learning_rate=0.06, max_iter=200, max_depth=3)
        mdl.fit(Xtr[m], ytr[m].astype(int))
        return mdl.predict_proba(np.nan_to_num(Xte))[:, 1]
    raise ValueError(kind)


def _bma_weights(P, y):
    """Brier-optimal non-negative weights over members (simple projected-gradient), interpretable."""
    n, M = P.shape
    w = np.full(M, 1.0 / M)
    for _ in range(200):
        pred = P @ w
        grad = 2 * P.T @ (pred - y) / n
        w = w - 0.5 * grad
        w = np.clip(w, 0, None)
        s = w.sum()
        w = w / s if s > 0 else np.full(M, 1.0 / M)
    return w


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)

    rows = []
    for b in blocks:
        w_fb, _ = mv.family_balanced_weights(members, mask[b])
        act = mask[b]
        yA = {k: [] for k in ("y", "count", "stack_logit", "stack_hgb", "bma")}
        for y in years:
            tr = [t for t in years if t != y]
            r = data[y][b]
            # COUNTING baseline: family-balanced vote fraction, in-CV isotonic recal
            fr_te = r["V"].astype(float) @ w_fb
            fr_tr = np.concatenate([data[t][b]["V"].astype(float) @ w_fb for t in tr])
            ytr = np.concatenate([data[t][b]["y"] for t in tr])
            mm = np.isfinite(fr_tr) & np.isfinite(ytr)
            iso = IsotonicRegression(out_of_bounds="clip").fit(fr_tr[mm], ytr[mm].astype(int))
            yA["count"].append(iso.transform(fr_te))
            # STACK base features = ACTIVE member probabilities (OOF). Meta-learners need far
            # fewer than ~1M rows; subsample the meta-training set (seeded) for feasibility.
            Ptr = np.concatenate([data[t][b]["P"][:, act] for t in tr])
            Pte = r["P"][:, act]
            rng = np.random.default_rng(12345 + hash(b) % 1000 + y)
            cap = 150000
            sel = np.arange(Ptr.shape[0]) if Ptr.shape[0] <= cap else rng.choice(Ptr.shape[0], cap, replace=False)
            Ptr_s, ytr_s = Ptr[sel], ytr[sel]
            for kind in ("stack_logit", "stack_hgb"):
                yA[kind].append(_fit_apply(kind, Ptr_s, ytr_s, Pte))
            # BMA on active member probs (subsampled)
            mb = np.isfinite(ytr_s) & np.all(np.isfinite(Ptr_s), axis=1)
            wbma = _bma_weights(np.nan_to_num(Ptr_s[mb]), ytr_s[mb].astype(float))
            yA["bma"].append(np.nan_to_num(Pte) @ wbma)
            yA["y"].append(r["y"])
        y = np.concatenate(yA["y"]); clim = np.full_like(y, np.nanmean(y), float)
        for model in ("count", "stack_logit", "stack_hgb", "bma"):
            p = np.concatenate(yA[model])
            # in-CV Peirce-opt threshold pooled (already OOS); use a single pooled threshold for the
            # binary summary (consistent across models) then contingency scores
            t0, _ = vf.best_threshold(y, p, "Peirce")
            s = vf.binary_scores(y, p, t0)
            rows.append(dict(block=b, model=model, bss_vs_clim=round(vf.brier_skill_score(y, p, clim), 4),
                             auc=round(vf.auc(y, p), 4), SEDI=round(s["SEDI"], 4),
                             Peirce=round(s["Peirce"], 4), thr=round(t0, 3)))
        c = next(x for x in rows if x["block"] == b and x["model"] == "count")
        best = max((x for x in rows if x["block"] == b and x["model"] != "count"),
                   key=lambda x: x["SEDI"])
        log.info("%s: count SEDI=%.3f | best stack (%s) SEDI=%.3f | delta=%+.3f",
                 b, c["SEDI"], best["model"], best["SEDI"], best["SEDI"] - c["SEDI"])

    df = pd.DataFrame(rows)
    df.to_csv(os.path.join(tdir, "v3_stack_vs_count.csv"), index=False)

    import matplotlib; matplotlib.use("Agg"); import matplotlib.pyplot as plt
    piv = df.pivot(index="block", columns="model", values="SEDI").reindex(blocks)
    ax = piv.plot(marker="o", figsize=(9, 5))
    ax.set_ylabel("SEDI (pooled OOS)"); ax.set_xlabel("lead block")
    ax.set_title("Stage-23: stacking vs counting (family-balanced vote fraction) — SEDI by lead")
    ax.grid(alpha=.3)
    plt.tight_layout(); plt.savefig(os.path.join(fdir, "v3_stack_vs_count.png"), dpi=130)
    print(df.to_string(index=False))


if __name__ == "__main__":
    main()