"""Calibration model families for single-member MOS.

Each family maps block-aggregated deterministic-forecast features (see features.py) to a
probability P(target in block). All expose a common interface:  fit(X, y) -> self ;
predict_proba(X) -> 1-D array of P(y=1).  NaN rows in (X, y) are dropped inside fit.

Families:
  raw        non-trainable reference: uses the forecast block exceedance fraction (f_exceed)
             directly as the probability -- the "uncalibrated forecast" baseline.
  climat     non-trainable reference: predicts the training base rate per lead block.
  logistic   standardized logistic regression (statistical baseline).
  hgb        histogram gradient boosting classifier (ML; log-loss), the ML workhorse.
  rf         random forest classifier (ML alternative).

The priority ML families named in the brief (quantile regression forests, gradient boosting
with pinball loss, distributional networks) act on CONTINUOUS Tx; they are added in the
continuous-MOS module. For the probability product the classifiers above are the direct
single-member analogue and are what the decision threshold consumes.
"""
from __future__ import annotations
import os
import numpy as np
from sklearn.linear_model import LogisticRegression, Ridge

# Parallelism for tree libraries. Under SLURM we pack several folds per node, so the submit
# script sets EHD_NJOBS to the per-task cpu count to stop xgboost/lightgbm/RF from grabbing
# every core on the node. Default -1 (use all) for interactive single-fold runs.
_NJOBS = int(os.environ.get("EHD_NJOBS", "-1"))
from sklearn.preprocessing import StandardScaler
from sklearn.pipeline import make_pipeline
from sklearn.ensemble import HistGradientBoostingClassifier, RandomForestClassifier
from sklearn.ensemble import HistGradientBoostingRegressor
from sklearn.neural_network import MLPClassifier, MLPRegressor
from sklearn.isotonic import IsotonicRegression

F_EXCEED_COL = 3   # index of f_exceed in features.FEATURE_NAMES


class _Base:
    trainable = True
    def fit(self, X, y):
        return self
    def predict_proba(self, X):
        raise NotImplementedError


class RawForecast(_Base):
    """Uncalibrated reference: P = forecast block exceedance fraction (clipped to (eps,1-eps))."""
    trainable = False
    def predict_proba(self, X):
        p = X[:, F_EXCEED_COL].astype(float)
        return np.clip(p, 1e-3, 1 - 1e-3)


class Climatology(_Base):
    """Reference: constant training base rate."""
    trainable = False
    def fit(self, X, y):
        m = np.isfinite(y)
        self.rate_ = float(np.mean(y[m])) if m.any() else 0.0
        return self
    def predict_proba(self, X):
        return np.full(X.shape[0], getattr(self, "rate_", 0.0))


class Logistic(_Base):
    def __init__(self, C=1.0):
        self.model = make_pipeline(StandardScaler(),
                                   LogisticRegression(C=C, max_iter=1000))
    def fit(self, X, y):
        m = np.isfinite(y) & np.all(np.isfinite(X), axis=1)
        self.model.fit(X[m], y[m].astype(int)); return self
    def predict_proba(self, X):
        return self.model.predict_proba(X)[:, 1]


class HGB(_Base):
    def __init__(self, max_depth=None, learning_rate=0.06, max_iter=300, l2=0.0,
                 min_samples_leaf=200):
        self.model = HistGradientBoostingClassifier(
            max_depth=max_depth, learning_rate=learning_rate, max_iter=max_iter,
            l2_regularization=l2, min_samples_leaf=min_samples_leaf,
            early_stopping=True, validation_fraction=0.1, random_state=0)
    def fit(self, X, y):
        m = np.isfinite(y) & np.all(np.isfinite(X), axis=1)
        self.model.fit(X[m], y[m].astype(int)); return self
    def predict_proba(self, X):
        return self.model.predict_proba(X)[:, 1]


class RF(_Base):
    def __init__(self, n_estimators=300, max_depth=None, min_samples_leaf=100):
        self.model = RandomForestClassifier(
            n_estimators=n_estimators, max_depth=max_depth,
            min_samples_leaf=min_samples_leaf, n_jobs=_NJOBS, random_state=0)
    def fit(self, X, y):
        m = np.isfinite(y) & np.all(np.isfinite(X), axis=1)
        self.model.fit(X[m], y[m].astype(int)); return self
    def predict_proba(self, X):
        return self.model.predict_proba(X)[:, 1]


class MLP(_Base):
    """ANN classifier (multilayer perceptron) on standardized features."""
    def __init__(self, hidden=(64, 32), alpha=1e-3, max_iter=80):
        self.model = make_pipeline(
            StandardScaler(),
            MLPClassifier(hidden_layer_sizes=hidden, alpha=alpha, max_iter=max_iter,
                          early_stopping=True, n_iter_no_change=6, validation_fraction=0.1,
                          random_state=0))
    def fit(self, X, y):
        m = np.isfinite(y) & np.all(np.isfinite(X), axis=1)
        self.model.fit(X[m], y[m].astype(int)); return self
    def predict_proba(self, X):
        return self.model.predict_proba(X)[:, 1]


class XGB(_Base):
    """XGBoost gradient-boosting classifier (challenger to HGB; boosting diversity).

    NaN rows are handled natively by XGBoost, but we drop NaN-y rows to match the other
    families. Deterministic (seed=0), histogram tree method for speed on tabular predictors.
    """
    def __init__(self, n_estimators=400, max_depth=4, learning_rate=0.05,
                 subsample=0.8, colsample_bytree=0.8, reg_lambda=1.0, min_child_weight=20):
        from xgboost import XGBClassifier
        self.model = XGBClassifier(
            n_estimators=n_estimators, max_depth=max_depth, learning_rate=learning_rate,
            subsample=subsample, colsample_bytree=colsample_bytree, reg_lambda=reg_lambda,
            min_child_weight=min_child_weight, tree_method="hist", eval_metric="logloss",
            n_jobs=_NJOBS, random_state=0)
    def fit(self, X, y):
        m = np.isfinite(y)
        self.model.fit(X[m], y[m].astype(int)); return self
    def predict_proba(self, X):
        return self.model.predict_proba(X)[:, 1]


class LGBM(_Base):
    """LightGBM gradient-boosting classifier (fast challenger; boosting diversity)."""
    def __init__(self, n_estimators=500, num_leaves=31, learning_rate=0.05,
                 subsample=0.8, colsample_bytree=0.8, reg_lambda=1.0, min_child_samples=50):
        from lightgbm import LGBMClassifier
        self.model = LGBMClassifier(
            n_estimators=n_estimators, num_leaves=num_leaves, learning_rate=learning_rate,
            subsample=subsample, subsample_freq=1, colsample_bytree=colsample_bytree,
            reg_lambda=reg_lambda, min_child_samples=min_child_samples,
            n_jobs=_NJOBS, random_state=0, verbosity=-1)
    def fit(self, X, y):
        m = np.isfinite(y)
        self.model.fit(X[m], y[m].astype(int)); return self
    def predict_proba(self, X):
        return self.model.predict_proba(X)[:, 1]


# ---------------------------------------------------------------- severity regressors (v2)
class SevRegressor(_Base):
    """Predictand = severity Sev = Tx - Threshold (block max), NOT the binary label.

    Fits a regressor on continuous `sev`, then maps predicted sev -> P(yEHDE) with an
    isotonic calibration fit on TRAIN (predicted_sev, y). By construction sev>=0 <=> EHD, so
    the isotonic map is monotone and turns the severity forecast into a probability. Interface:
      fit(X, sev, y) -> self ;  predict_proba(X) -> P(y=1).
    The v2 runner calls fit(X, sev, y); the leakage rule is enforced by the caller (train-only).
    """
    is_sev = True
    def __init__(self, base):
        self.base = base
    def fit(self, X, sev, y):
        m = np.all(np.isfinite(X), axis=1) & np.isfinite(sev)
        self.base.fit(X[m], sev[m])
        s = self.base.predict(X)
        m2 = np.isfinite(s) & np.isfinite(y)
        self.iso_ = IsotonicRegression(out_of_bounds="clip").fit(s[m2], y[m2].astype(int))
        return self
    def predict_proba(self, X):
        return self.iso_.transform(self.base.predict(X))


def _ridge_sev(alpha=1.0):
    return SevRegressor(make_pipeline(StandardScaler(), Ridge(alpha=alpha)))

def _hgb_sev(learning_rate=0.06, max_iter=300, min_samples_leaf=200, l2=0.0):
    return SevRegressor(HistGradientBoostingRegressor(
        learning_rate=learning_rate, max_iter=max_iter, min_samples_leaf=min_samples_leaf,
        l2_regularization=l2, early_stopping=True, validation_fraction=0.1, random_state=0))

def _mlp_sev(hidden=(64, 32), alpha=1e-3, max_iter=80):
    return SevRegressor(make_pipeline(StandardScaler(), MLPRegressor(
        hidden_layer_sizes=hidden, alpha=alpha, max_iter=max_iter, early_stopping=True,
        n_iter_no_change=6, validation_fraction=0.1, random_state=0)))

def _xgb_sev(n_estimators=400, max_depth=4, learning_rate=0.05):
    from xgboost import XGBRegressor
    return SevRegressor(XGBRegressor(
        n_estimators=n_estimators, max_depth=max_depth, learning_rate=learning_rate,
        subsample=0.8, colsample_bytree=0.8, reg_lambda=1.0, min_child_weight=20,
        tree_method="hist", n_jobs=_NJOBS, random_state=0))

def _lgbm_sev(n_estimators=500, num_leaves=31, learning_rate=0.05):
    from lightgbm import LGBMRegressor
    return SevRegressor(LGBMRegressor(
        n_estimators=n_estimators, num_leaves=num_leaves, learning_rate=learning_rate,
        subsample=0.8, subsample_freq=1, colsample_bytree=0.8, reg_lambda=1.0,
        min_child_samples=50, n_jobs=_NJOBS, random_state=0, verbosity=-1))


REGISTRY = {"raw": RawForecast, "climat": Climatology,
            "logistic": Logistic, "hgb": HGB, "rf": RF, "mlp": MLP,
            "xgb": XGB, "lgbm": LGBM,
            "ridge_sev": _ridge_sev, "hgb_sev": _hgb_sev, "mlp_sev": _mlp_sev,
            "xgb_sev": _xgb_sev, "lgbm_sev": _lgbm_sev}

# families whose fit signature is fit(X, sev, y) rather than fit(X, y)
SEV_FAMILIES = {"ridge_sev", "hgb_sev", "mlp_sev", "xgb_sev", "lgbm_sev"}


def is_sev_family(name):
    return name in SEV_FAMILIES


def make_model(name, **kw):
    return REGISTRY[name](**kw)