"""Stage-2 item-2 expansion: class-imbalance-aware member configurations.

These are NEW configurations (registry v3) of families that already earned their place in
Stage 1, addressing the rare-event imbalance directly with class-weighted losses. They present
the same interface as `ehd.calibration` classifiers (`fit(X, y)`, `predict_proba(X) -> P(event)`),
so the Stage-1 training/recalibration machinery is reused unchanged. NaNs are masked on fit and
zero-filled on predict, matching the zoo.
"""
from __future__ import annotations
import numpy as np
from sklearn.linear_model import LogisticRegression
from sklearn.ensemble import RandomForestClassifier
from sklearn.preprocessing import StandardScaler
from sklearn.pipeline import make_pipeline


class _ProbaClf:
    def __init__(self, model):
        self.model = model

    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(np.nan_to_num(X, nan=0.0))[:, 1]


def balanced_logistic(C=1.0):
    """Standardised logistic regression with class_weight='balanced'."""
    return _ProbaClf(make_pipeline(
        StandardScaler(),
        LogisticRegression(C=C, max_iter=1000, class_weight="balanced")))


def balanced_rf(n_estimators=300, min_samples_leaf=100):
    """Random forest with class_weight='balanced_subsample' for the rare positive class."""
    return _ProbaClf(RandomForestClassifier(
        n_estimators=n_estimators, min_samples_leaf=min_samples_leaf,
        class_weight="balanced_subsample", n_jobs=-1, random_state=0))


def make_balanced(family):
    return {"logistic": balanced_logistic, "rf": balanced_rf}[family]()