"""Two-level (cascade) classifier for EHD/EHDE — v2 method.

Request (user 2026-09-14): screen out the abundant, trivially-classified no-event samples
before training the event models, so the hard classifier is not swamped by easy negatives.

  LV1  a screening classifier forecasts P(yEHDE) (yEHDE = >=1 EHD in the block). Its
       probabilities are isotonic-recalibrated on TRAIN so the cutoff has a calibrated
       meaning. A sample is "nearly certain nEHDE" when P(yEHDE) <= 1 - cutoff; those are
       REMOVED from the LV2 pool and simply keep their LV1 probability (near 0).
  LV2  the event model (any zoo family, incl. severity regressors) is trained ONLY on the
       surviving TRAIN samples and scores the surviving TEST samples. Screened TEST samples
       take the LV1 probability.

`cutoff` sweeps from 0.90 upward per the user's instruction; `cutoff = 1.0` screens nothing
(1 - cutoff = 0, and screening is `p1 <= 1-cutoff`), giving the single-level baseline for free.

LEAKAGE RULE: LV1 and LV2 are fit on TRAIN arrays only; the caller passes train/test already
split by held-out year. Nothing here touches the held-out year.
"""
from __future__ import annotations
import numpy as np
from sklearn.isotonic import IsotonicRegression
from . import calibration as cal


def _recalibrate(p_tr, y_tr, *arrays):
    """Isotonic map fit on (p_tr, y_tr); returns transformed p_tr plus each extra array."""
    m = np.isfinite(p_tr) & np.isfinite(y_tr)
    iso = IsotonicRegression(out_of_bounds="clip").fit(p_tr[m], y_tr[m].astype(int))
    return (iso.transform(p_tr), *(iso.transform(a) for a in arrays))


def screen_lv1(Xtr, ytr, Xte, lv1_family="hgb"):
    """Fit the LV1 screening classifier; return calibrated P(yEHDE) on train and test."""
    lv1 = cal.make_model(lv1_family).fit(Xtr, ytr)
    r_tr = lv1.predict_proba(Xtr)
    r_te = lv1.predict_proba(Xte)
    p1_tr, p1_te = _recalibrate(r_tr, ytr, r_te)
    return p1_tr, p1_te


def run_cascade(Xtr, ytr, sevtr, Xte, family, cutoff,
                p1_tr, p1_te, recal_lv2=True):
    """One LV2 family at one cutoff, reusing precomputed LV1 probabilities.

    Returns (p_te, p_tr, info) where p_* are the FINAL combined probabilities (screened
    samples keep LV1, survivors take LV2) on test and train. p_tr is returned so the caller
    can fit the per-region decision threshold on TRAIN (leakage-safe).
    """
    thr_screen = 1.0 - cutoff
    keep_tr = p1_tr > thr_screen
    keep_te = p1_te > thr_screen
    # guard: if screening leaves too few survivors or one class, fall back to no screening
    ys = ytr[keep_tr]
    if keep_tr.sum() < 500 or np.nansum(ys == 1) < 20 or np.nansum(ys == 0) < 20:
        keep_tr = np.ones_like(keep_tr); keep_te = np.ones_like(keep_te)

    if cal.is_sev_family(family):
        lv2 = cal.make_model(family).fit(Xtr[keep_tr], sevtr[keep_tr], ytr[keep_tr])
        p2_tr_full = lv2.predict_proba(Xtr)      # sev regressors self-calibrate (isotonic)
        p2_te = lv2.predict_proba(Xte)
    else:
        lv2 = cal.make_model(family).fit(Xtr[keep_tr], ytr[keep_tr])
        p2_tr_full = lv2.predict_proba(Xtr)
        p2_te = lv2.predict_proba(Xte)
        if recal_lv2:
            # recalibrate on the survivor training set (where LV2 was actually fit)
            p2_tr_s, = _recalibrate(p2_tr_full[keep_tr], ytr[keep_tr])
            # apply the same map to all train/test survivor scores
            m = np.isfinite(p2_tr_full[keep_tr]) & np.isfinite(ytr[keep_tr])
            iso = IsotonicRegression(out_of_bounds="clip").fit(
                p2_tr_full[keep_tr][m], ytr[keep_tr][m].astype(int))
            p2_tr_full = iso.transform(p2_tr_full)
            p2_te = iso.transform(p2_te)

    p_tr = np.where(keep_tr, p2_tr_full, p1_tr)
    p_te = np.where(keep_te, p2_te, p1_te)
    info = dict(screened_frac_tr=float(1.0 - keep_tr.mean()),
                screened_frac_te=float(1.0 - keep_te.mean()))
    return p_te, p_tr, info