Blame for src/batteryswapai/competition_model.py

from __future__ import annotations

from dataclasses import dataclass

import numpy as np
import pandas as pd
from sklearn.ensemble import (
    HistGradientBoostingClassifier,
    HistGradientBoostingRegressor,
    RandomForestClassifier,
)
from sklearn.impute import SimpleImputer
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import GroupKFold
from scipy.stats import rankdata

from .competition_features import (
    CONTEXT_COLUMNS,
    MIN_TRAJECTORY_COVERAGE,
    TRAJECTORY_COLUMNS,
    TRAJECTORY_REFERENCE_OFFSET,
    TRAJECTORY_COVERAGE_COLUMN,
    TRAJECTORY_LAGS,
    add_context_columns,
    numeric_feature_columns,
)


def _logit(values: np.ndarray) -> np.ndarray:
    clipped = np.clip(values, 1e-5, 1.0 - 1e-5)
    return np.log(clipped / (1.0 - clipped))


@dataclass
class FrozenLogisticCalibrator:
    """Minimal binary logistic inference independent of sklearn pickle internals."""

    coefficients: np.ndarray
    intercept: float

    @classmethod
    def from_estimator(cls, estimator: LogisticRegression) -> "FrozenLogisticCalibrator":
        return cls(
            coefficients=np.asarray(estimator.coef_[0], dtype=float).copy(),
            intercept=float(estimator.intercept_[0]),
        )

    def predict_proba(self, values: np.ndarray) -> np.ndarray:
        scores = np.asarray(values, dtype=float) @ self.coefficients + self.intercept
        positive = np.empty_like(scores, dtype=float)
        nonnegative = scores >= 0.0
        positive[nonnegative] = 1.0 / (1.0 + np.exp(-scores[nonnegative]))
        exp_scores = np.exp(scores[~nonnegative])
        positive[~nonnegative] = exp_scores / (1.0 + exp_scores)
        return np.column_stack([1.0 - positive, positive])


def _clean(frame: pd.DataFrame, columns: list[str]) -> pd.DataFrame:
    return frame.reindex(columns=columns).replace([np.inf, -np.inf], np.nan).astype(float)


def _derived_features(base: pd.DataFrame, kind: str) -> pd.DataFrame:
    derived: dict[str, pd.Series] = {}
    if kind in {"physics", "all"}:
        for level in ("voltage_mean_3d", "voltage_mean_7d", "voltage_max_3d"):
            for window in (30, 60, 90, 180):
                rate = np.maximum(-base[f"voltage_slope_{window}d"], 1e-5)
                derived[f"knee_{level}_{window}"] = np.clip(
                    (base[level] - 2.40) / rate,
                    -100.0,
                    1000.0,
                )
    if kind == "all":
        for short, long in ((3, 14), (7, 30), (14, 60), (30, 90), (30, 180), (60, 365)):
            derived[f"slope_accel_{short}_{long}"] = (
                base[f"voltage_slope_{short}d"] - base[f"voltage_slope_{long}d"]
            )
            derived[f"mean_drop_{short}_{long}"] = (
                base[f"voltage_mean_{short}d"] - base[f"voltage_mean_{long}d"]
            )
        for window in (3, 7, 14, 30, 60, 90):
            derived[f"volt_temp_{window}"] = base[f"voltage_mean_{window}d"] / (
                base[f"temperature_mean_{window}d"] + 273.15
            )
    return pd.concat([base, pd.DataFrame(derived, index=base.index)], axis=1)


HAZARD_EDGES = (0, 14, 28, 42, 70, 120, 200, 365)
FORECAST_HORIZONS = (30, 60, 90)
FORECAST_THRESHOLDS = (2.50, 2.45)
TRAJECTORY_BLEND_WEIGHT = 0.25


def _trajectory_features(bins: np.ndarray) -> pd.DataFrame:
    """Shape of the degradation trajectory: levels, drops, steps and summaries."""

    current = bins[:, 0]
    highest = np.nanmax(bins, axis=1)
    lowest = np.nanmin(bins, axis=1)
    spread = np.nanstd(bins, axis=1)
    columns: dict[str, np.ndarray] = {}
    for position, lag in enumerate(TRAJECTORY_LAGS):
        columns[f"v_lag{lag}"] = bins[:, position]
    for position, lag in enumerate(TRAJECTORY_LAGS[1:], start=1):
        columns[f"drop{lag}"] = current - bins[:, position]
        columns[f"ndrop{lag}"] = (current - bins[:, position]) / np.where(
            spread > 1e-6, spread, np.nan
        )
    steps = np.diff(bins, axis=1)
    for position in range(steps.shape[1]):
        columns[f"step{position}"] = -steps[:, position]
    columns.update(
        own_max=highest,
        own_min=lowest,
        own_std=spread,
        cur_minus_max=current - highest,
        cur_over_max=current / highest,
        pos=(current - lowest) / np.where(highest - lowest > 1e-6, highest - lowest, np.nan),
        fb50=np.nanmean(bins < 2.50, axis=1),
        fb45=np.nanmean(bins < 2.45, axis=1),
        mono=np.nanmean(steps < 0, axis=1),
    )
    return pd.DataFrame(columns)


def _snapshot_bins(snapshot: pd.DataFrame) -> np.ndarray | None:
    if not set(TRAJECTORY_COLUMNS).issubset(snapshot.columns):
        return None
    return snapshot.reindex(columns=list(TRAJECTORY_COLUMNS)).to_numpy(dtype=float)


def _dense_transitions(matrix: dict, buildings: dict[str, str]):
    """Every observed degradation transition, not just the terminal failures."""

    grid = matrix["grid"]
    first = matrix["first"]
    devices = list(matrix["index"])
    rows, cuts = [], []
    for name, position in matrix["index"].items():
        begin = int(first[position])
        if begin >= grid.shape[1]:
            continue
        end = int(np.flatnonzero(np.isfinite(grid[position]))[-1])
        start = begin + 180
        for cut in range(start, end + 1, 14):
            rows.append(position)
            cuts.append(cut)
    rows = np.asarray(rows, dtype=np.int64)
    cuts = np.asarray(cuts, dtype=np.int64)
    lags = np.asarray(TRAJECTORY_LAGS, dtype=np.int64) + TRAJECTORY_REFERENCE_OFFSET
    columns = cuts[:, None] - lags[None, :]
    safe = np.clip(columns, 0, grid.shape[1] - 1)
    values = grid[rows[:, None], safe]
    usable = (columns >= 0) & (columns >= first[rows][:, None])
    bins = np.where(usable, values, np.nan)
    names = np.array([devices[position] for position in rows])
    where = np.array([buildings.get(name, "") for name in names])
    return rows, cuts, bins, names, where


def _forecast_targets(matrix: dict, rows, cuts, eol_day: dict[str, float], names):
    """Threshold crossings and future voltage, read only from observed history."""

    grid = matrix["grid"]
    last = matrix["observed_last"]
    ends_at = np.array([eol_day.get(name, np.nan) for name in names], dtype=float)
    targets: dict[tuple, np.ndarray] = {}
    for horizon in FORECAST_HORIZONS:
        finish = cuts + horizon
        reachable = finish <= last[rows]
        died = np.isfinite(ends_at) & (ends_at <= finish)
        lowest = np.full(len(rows), np.nan)
        for position in range(len(rows)):
            begin = cuts[position] + 1
            stop = min(int(finish[position]), int(last[rows[position]]))
            if stop >= begin:
                lowest[position] = np.nanmin(grid[rows[position], begin : stop + 1])
        for threshold in FORECAST_THRESHOLDS:
            label = np.where(
                died, 1.0, np.where(reachable, (lowest <= threshold).astype(float), np.nan)
            )
            targets[(horizon, threshold)] = np.where(~reachable & ~died, np.nan, label)
        future = np.where(reachable, grid[rows, np.clip(finish, 0, grid.shape[1] - 1)], np.nan)
        targets[(horizon, "voltage")] = np.where(died, 2.30, future)
    return targets


def _forecaster(random_state: int, regression: bool):
    common = dict(learning_rate=0.06, max_iter=200, max_leaf_nodes=15, min_samples_leaf=40,
                  l2_regularization=5.0, early_stopping=True, validation_fraction=0.1,
                  n_iter_no_change=15, random_state=random_state)
    if regression:
        return HistGradientBoostingRegressor(**common)
    return HistGradientBoostingClassifier(loss="log_loss", **common)


def _fit_forecasters(matrix, buildings, eol_day, random_state,
                     snapshot_features, snapshot_buildings):
    """Learn degradation dynamics from dense transitions; score snapshots by held-out building."""

    rows, cuts, bins, names, where = _dense_transitions(matrix, buildings)
    if len(rows) < 500:
        return None, None
    features = _trajectory_features(bins)
    targets = _forecast_targets(matrix, rows, cuts, eol_day, names)
    unique = np.array(sorted(set(where) - {""}))
    if unique.size < 3:
        return None, None
    assignment = {name: position % 3 for position, name in enumerate(unique)}
    dense_fold = np.array([assignment.get(name, -1) for name in where])
    snap_fold = np.array([assignment.get(str(name), -1) for name in snapshot_buildings])

    fitted: dict[str, object] = {}
    holdout: dict[str, np.ndarray] = {}
    for (horizon, kind), values in targets.items():
        regression = kind == "voltage"
        name = f"v{horizon}" if regression else f"p{horizon}_{kind}"
        usable = np.isfinite(values)
        if usable.sum() < 200:
            continue
        estimator = _forecaster(random_state, regression)
        estimator.fit(features[usable], values[usable])
        fitted[name] = estimator
        out = np.full(len(snapshot_features), np.nan)
        for group in range(3):
            train = usable & (dense_fold != group)
            target_rows = np.flatnonzero(snap_fold == group)
            if train.sum() < 100 or target_rows.size == 0:
                continue
            fold_model = _forecaster(random_state + group + 1, regression)
            fold_model.fit(features[train], values[train])
            block = snapshot_features.iloc[target_rows]
            out[target_rows] = (
                fold_model.predict(block) if regression
                else fold_model.predict_proba(block)[:, 1]
            )
        holdout[name] = out
    if not fitted:
        return None, None
    return fitted, pd.DataFrame(holdout)


def _forecast_features(forecasters, features: pd.DataFrame) -> pd.DataFrame:
    out: dict[str, np.ndarray] = {}
    for name, estimator in forecasters.items():
        if name.startswith("v"):
            out[name] = estimator.predict(features)
        else:
            out[name] = estimator.predict_proba(features)[:, 1]
    return pd.DataFrame(out, index=features.index)



def _context_features(snapshot: pd.DataFrame) -> pd.DataFrame:
    """Scenario-level context, already attached by scenario_snapshot."""

    missing = [name for name in CONTEXT_COLUMNS if name not in snapshot.columns]
    if missing:
        snapshot = add_context_columns(snapshot.copy())
    return snapshot.reindex(columns=list(CONTEXT_COLUMNS)).astype(float)


def _context_estimators(random_state: int) -> dict[str, object]:
    estimators: dict[str, object] = {}
    for seed in range(3):
        estimators[f"hist_{seed}"] = HistGradientBoostingClassifier(
            loss="log_loss", learning_rate=0.05, max_iter=300, max_leaf_nodes=15,
            min_samples_leaf=30, l2_regularization=5.0, early_stopping=True,
            validation_fraction=0.12, n_iter_no_change=25,
            class_weight={0: 1.0, 1: 5.0}, random_state=random_state + seed,
        )
    estimators["forest"] = RandomForestClassifier(
        n_estimators=350, min_samples_leaf=5, max_features=0.7,
        class_weight="balanced_subsample", n_jobs=-1, random_state=random_state,
    )
    return estimators


def _fit_context(features: pd.DataFrame, target: np.ndarray, random_state: int):
    estimators = _context_estimators(random_state)
    imputer = SimpleImputer(strategy="median")
    imputed = imputer.fit_transform(features)
    for name, estimator in estimators.items():
        estimator.fit(imputed if name == "forest" else features, target)
    return estimators, imputer


def _context_risk(estimators, imputer, features: pd.DataFrame) -> np.ndarray:
    imputed = imputer.transform(features)
    parts = [
        estimator.predict_proba(imputed if name == "forest" else features)[:, 1]
        for name, estimator in estimators.items()
    ]
    return np.mean(parts, axis=0)


def _fit_hazard(base: pd.DataFrame, target_rul: pd.Series, observed: pd.Series, random_state: int):
    """Person-period model: one row per battery per interval it is still alive."""

    rul = target_rul.to_numpy(dtype=float)
    events = observed.to_numpy(dtype=bool)
    rows, bins, labels = [], [], []
    for index in range(len(HAZARD_EDGES) - 1):
        low, high = HAZARD_EDGES[index], HAZARD_EDGES[index + 1]
        alive = np.flatnonzero(rul >= low)
        if alive.size == 0:
            continue
        rows.append(alive)
        bins.append(np.full(alive.shape, index))
        labels.append((events[alive] & (rul[alive] < high)).astype(int))
    if not rows:
        return None
    positions = np.concatenate(rows)
    features = base.iloc[positions].copy()
    features["hazard_interval"] = np.concatenate(bins)
    target = np.concatenate(labels)
    if target.sum() < 20:
        return None
    estimator = HistGradientBoostingClassifier(
        loss="log_loss", learning_rate=0.05, max_iter=300, max_leaf_nodes=15,
        min_samples_leaf=40, l2_regularization=5.0, early_stopping=True,
        validation_fraction=0.12, n_iter_no_change=25, random_state=random_state,
    )
    estimator.fit(features, target)
    return estimator


def _hazard_risk(estimator, base: pd.DataFrame, horizon_days: int) -> np.ndarray:
    survival = np.ones(len(base), dtype=float)
    for index in range(len(HAZARD_EDGES) - 1):
        if HAZARD_EDGES[index + 1] > horizon_days:
            break
        features = base.copy()
        features["hazard_interval"] = index
        survival *= 1.0 - estimator.predict_proba(features)[:, 1]
    return 1.0 - survival


def _scenario_ranks(values: np.ndarray, groups: np.ndarray | None) -> np.ndarray:
    if groups is None:
        return rankdata(values) / max(len(values), 1)
    ranks = np.empty(len(values), dtype=float)
    for key in pd.unique(groups):
        mask = groups == key
        ranks[mask] = rankdata(values[mask]) / max(int(mask.sum()), 1)
    return ranks


def _event_estimators(random_state: int) -> dict[str, object]:
    common = dict(
        loss="log_loss",
        learning_rate=0.05,
        max_iter=240,
        early_stopping=True,
        validation_fraction=0.12,
        n_iter_no_change=25,
        random_state=random_state,
    )
    return {
        "weighted_hist": HistGradientBoostingClassifier(
            max_leaf_nodes=15,
            min_samples_leaf=35,
            l2_regularization=4.0,
            class_weight={0: 1.0, 1: 5.0},
            **common,
        ),
        "forest": RandomForestClassifier(
            n_estimators=350,
            min_samples_leaf=5,
            max_features=0.7,
            class_weight="balanced_subsample",
            n_jobs=-1,
            random_state=random_state,
        ),
        "all_hist": HistGradientBoostingClassifier(
            max_leaf_nodes=31,
            min_samples_leaf=25,
            l2_regularization=5.0,
            **common,
        ),
        "physics_hist": HistGradientBoostingClassifier(
            max_leaf_nodes=31,
            min_samples_leaf=25,
            l2_regularization=5.0,
            **common,
        ),
    }


def _fit_event_estimators(
    estimators: dict[str, object],
    base: pd.DataFrame,
    all_features: pd.DataFrame,
    physics_features: pd.DataFrame,
    target: np.ndarray,
    sample_weight: np.ndarray | None = None,
) -> tuple[dict[str, object], SimpleImputer]:
    imputer = SimpleImputer(strategy="median")
    forest_features = imputer.fit_transform(base)
    estimators["weighted_hist"].fit(base, target, sample_weight=sample_weight)
    estimators["forest"].fit(forest_features, target, sample_weight=sample_weight)
    estimators["all_hist"].fit(all_features, target, sample_weight=sample_weight)
    estimators["physics_hist"].fit(physics_features, target, sample_weight=sample_weight)
    return estimators, imputer


def _device_weights(devices: pd.Series) -> np.ndarray:
    """One vote per battery, so devices seen in many scenarios do not dominate."""

    counts = devices.groupby(devices).transform("size").to_numpy(dtype=float)
    weights = 1.0 / np.maximum(counts, 1.0)
    return weights * (len(weights) / weights.sum())


def _event_components(
    estimators: dict[str, object],
    imputer: SimpleImputer,
    base: pd.DataFrame,
    all_features: pd.DataFrame,
    physics_features: pd.DataFrame,
) -> np.ndarray:
    return np.column_stack(
        [
            estimators["weighted_hist"].predict_proba(base)[:, 1],
            estimators["forest"].predict_proba(imputer.transform(base))[:, 1],
            estimators["all_hist"].predict_proba(all_features)[:, 1],
            estimators["physics_hist"].predict_proba(physics_features)[:, 1],
        ]
    )


@dataclass
class EventTimeModel:
    event_estimators: dict[str, object]
    event_imputer: SimpleImputer
    event_calibrator: FrozenLogisticCalibrator
    rul_estimator: HistGradientBoostingRegressor
    survivor_rul_estimator: HistGradientBoostingRegressor | None
    feature_columns: list[str]
    all_feature_columns: list[str]
    physics_feature_columns: list[str]
    quantile: float
    horizon_days: int
    dataset_revision: str
    balanced: "EventTimeModel | None" = None
    context_estimators: dict | None = None
    context_imputer: SimpleImputer | None = None
    context_columns: list[str] | None = None
    hazard_estimator: object | None = None
    blend_weights: tuple[float, float, float, float] = (1.0, 0.0, 0.0, 0.0)
    forecasters: dict | None = None
    residual_model: object | None = None
    residual_columns: list[str] | None = None
    trajectory_weight: float = 0.0

    def _feature_sets(
        self, snapshot: pd.DataFrame
    ) -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]:
        base = _clean(snapshot, self.feature_columns)
        all_features = _clean(_derived_features(base, "all"), self.all_feature_columns)
        physics_features = _clean(
            _derived_features(base, "physics"), self.physics_feature_columns
        )
        return base, all_features, physics_features

    def _calibrated_risk(self, snapshot: pd.DataFrame) -> np.ndarray:
        base, all_features, physics_features = self._feature_sets(snapshot)
        components = _event_components(
            self.event_estimators,
            self.event_imputer,
            base,
            all_features,
            physics_features,
        )
        stacked = _logit(components)
        return np.asarray(self.event_calibrator.predict_proba(stacked)[:, 1])

    def predict_risk_components(self, snapshot: pd.DataFrame) -> dict[str, np.ndarray]:
        components = {"primary": self._calibrated_risk(snapshot)}
        if self.balanced is not None:
            components["balanced"] = self.balanced._calibrated_risk(snapshot)
        if self.context_estimators is not None:
            features = _clean(
                pd.concat(
                    [_clean(snapshot, self.feature_columns), _context_features(snapshot)],
                    axis=1,
                ),
                self.context_columns,
            )
            components["context"] = _context_risk(
                self.context_estimators, self.context_imputer, features
            )
        if self.hazard_estimator is not None:
            components["hazard"] = _hazard_risk(
                self.hazard_estimator, _clean(snapshot, self.feature_columns), self.horizon_days
            )
        return components

    def _residual_risk(self, snapshot: pd.DataFrame):
        """Degradation-dynamics view; None when the trajectory history is unusable."""

        if self.residual_model is None or self.forecasters is None:
            return None, None
        bins = _snapshot_bins(snapshot)
        if bins is None:
            return None, None
        trajectory = _trajectory_features(bins)
        forecast = _forecast_features(self.forecasters, trajectory)
        frame = pd.concat([forecast, trajectory], axis=1)
        frame = frame.reindex(columns=self.residual_columns).astype(float)
        risk = self.residual_model.predict_proba(frame)[:, 1]
        if TRAJECTORY_COVERAGE_COLUMN in snapshot.columns:
            coverage = pd.to_numeric(
                snapshot[TRAJECTORY_COVERAGE_COLUMN], errors="coerce"
            ).to_numpy(dtype=float)
        else:
            coverage = np.isfinite(bins).mean(axis=1)
        coverage = np.where(np.isfinite(coverage), coverage, 0.0)
        # A battery without enough history keeps its blend position rather than moving on noise.
        weight = np.where(coverage >= MIN_TRAJECTORY_COVERAGE, self.trajectory_weight, 0.0)
        return risk, weight

    def predict_event_risk(self, snapshot: pd.DataFrame) -> np.ndarray:
        groups = (
            snapshot["scenario"].astype(str).to_numpy()
            if "scenario" in snapshot.columns
            else None
        )
        primary = blend_risk_components(
            self.predict_risk_components(snapshot), self.blend_weights, groups
        )
        risk, weight = self._residual_risk(snapshot)
        if risk is None or not np.any(weight > 0.0):
            return primary
        return blend_trajectory(primary, risk, weight, groups)

    def predict_rul(self, snapshot: pd.DataFrame) -> np.ndarray:
        base = _clean(snapshot, self.feature_columns)
        return np.asarray(self.rul_estimator.predict(base), dtype=float)

    def predict_survivor_rul(self, snapshot: pd.DataFrame) -> np.ndarray | None:
        """Expected remaining life of batteries that outlive the planning window."""

        if self.survivor_rul_estimator is None:
            return None
        base = _clean(snapshot, self.feature_columns)
        return np.asarray(self.survivor_rul_estimator.predict(base), dtype=float)


def blend_risk_components(
    components: dict[str, np.ndarray],
    weights: tuple[float, float, float, float],
    groups: np.ndarray | None,
) -> np.ndarray:
    """Rank each view inside a scenario, then map back onto the calibrated scale."""

    order = ("primary", "balanced", "context", "hazard")
    primary = np.asarray(components["primary"], dtype=float)
    active = [
        (np.asarray(components[name], dtype=float), float(weights[index]))
        for index, name in enumerate(order)
        if name in components and float(weights[index]) > 0.0
    ]
    if len(active) <= 1:
        return primary
    blended = np.zeros(len(primary), dtype=float)
    total = 0.0
    for values, weight in active:
        blended += weight * _scenario_ranks(values, groups)
        total += weight
    blended /= max(total, 1e-9)
    return _map_to_scale(blended, primary, groups)


def blend_trajectory(primary, residual, weight, groups):
    """Nudge the ranking by the degradation view, keeping the calibrated scale."""

    primary = np.asarray(primary, dtype=float)
    residual = np.asarray(residual, dtype=float)
    weight = np.asarray(weight, dtype=float)
    residual = np.where(np.isfinite(residual), residual, np.nanmin(residual[np.isfinite(residual)])
                        if np.isfinite(residual).any() else 0.0)
    blended = np.empty(len(primary), dtype=float)
    keys = [None] if groups is None else list(pd.unique(groups))
    for key in keys:
        mask = slice(None) if key is None else (groups == key)
        size = len(primary[mask])
        first = rankdata(primary[mask]) / max(size, 1)
        second = rankdata(residual[mask]) / max(size, 1)
        local = weight[mask]
        blended[mask] = (1.0 - local) * first + local * second
    return _map_to_scale(blended, primary, groups)


def _map_to_scale(blended: np.ndarray, reference: np.ndarray, groups) -> np.ndarray:
    """Keep the calibrated marginal distribution and take only the new ordering."""

    mapped = np.empty(len(blended), dtype=float)
    keys = [None] if groups is None else list(pd.unique(groups))
    for key in keys:
        mask = slice(None) if key is None else (groups == key)
        scores = blended[mask]
        scale = np.sort(np.asarray(reference)[mask])
        order = rankdata(scores, method="ordinal").astype(int) - 1
        mapped[mask] = scale[order]
    return mapped


def _fit_trajectory_stack(work, trajectory, eol_times, event_target, random_state):
    """Dense degradation forecasting plus the temporal-bin view, as one residual ranker."""

    bins = _snapshot_bins(work)
    if bins is None:
        return None, None, None, 0.0
    snapshot_features = _trajectory_features(bins)
    buildings = dict(zip(work["battery"].astype(str), work["building"].astype(str)))
    eol_day: dict[str, float] = {}
    if eol_times is not None:
        origin = trajectory["origin"]
        for name, when in eol_times.items():
            if pd.notna(when):
                eol_day[str(name)] = float((pd.Timestamp(when).normalize() - origin).days)
    forecasters, holdout = _fit_forecasters(
        trajectory, buildings, eol_day, random_state,
        snapshot_features, work["building"].astype(str).to_numpy(),
    )
    if forecasters is None:
        return None, None, None, 0.0
    frame = pd.concat([holdout.reset_index(drop=True), snapshot_features.reset_index(drop=True)], axis=1)
    columns = list(frame.columns)
    residual = HistGradientBoostingClassifier(
        loss="log_loss", learning_rate=0.05, max_iter=300, max_leaf_nodes=15,
        min_samples_leaf=30, l2_regularization=5.0, early_stopping=True,
        validation_fraction=0.12, n_iter_no_change=25, class_weight={0: 1.0, 1: 5.0},
        random_state=random_state,
    )
    residual.fit(frame.astype(float), event_target)
    return forecasters, residual, columns, TRAJECTORY_BLEND_WEIGHT


def _fit_survivor_estimator(
    base: pd.DataFrame,
    target_rul: pd.Series,
    observed: pd.Series,
    horizon_days: int,
    maximum_days: float,
    random_state: int,
) -> HistGradientBoostingRegressor | None:
    """Median remaining life among batteries observed to outlive the window."""

    mask = observed & target_rul.gt(float(horizon_days)) & target_rul.le(float(maximum_days))
    if int(mask.sum()) < 50:
        return None
    estimator = HistGradientBoostingRegressor(
        loss="quantile",
        quantile=0.5,
        learning_rate=0.05,
        max_iter=200,
        max_leaf_nodes=15,
        min_samples_leaf=40,
        l2_regularization=4.0,
        early_stopping=True,
        validation_fraction=0.15,
        n_iter_no_change=20,
        random_state=random_state,
    )
    estimator.fit(base.loc[mask], target_rul.loc[mask])
    return estimator


def fit_event_time_model(
    training_snapshots: pd.DataFrame,
    *,
    quantile: float = 0.05,
    horizon_days: int = 42,
    rul_training_max_days: float = 90.0,
    survivor_training_max_days: float = 400.0,
    dataset_revision: str = "unknown",
    random_state: int = 2026,
    calibration_folds: int = 5,
    device_balanced: bool = False,
    with_auxiliary: bool = True,
    trajectory: dict | None = None,
    eol_times: pd.Series | None = None,
    blend_weights: tuple[float, float, float, float] = (0.55, 0.15, 0.15, 0.15),
    max_iter: int | None = None,
) -> EventTimeModel:
    del max_iter  # retained for compatibility with earlier validation commands
    work = training_snapshots.dropna(subset=["target_rul_days"]).copy()
    columns = numeric_feature_columns(work)
    base = _clean(work, columns)
    all_features = _derived_features(base, "all")
    physics_features = _derived_features(base, "physics")
    all_columns = list(all_features.columns)
    physics_columns = list(physics_features.columns)
    target_rul = pd.to_numeric(work["target_rul_days"], errors="coerce").astype(float)
    event_target = (
        work["event_observed"].astype(bool)
        & target_rul.between(0.0, float(horizon_days))
    ).astype("int8").to_numpy()

    devices = work["battery"].astype(str)
    groups = devices.to_numpy()
    weights = _device_weights(devices) if device_balanced else None
    folds = min(int(calibration_folds), len(np.unique(groups)))
    if folds < 2:
        raise ValueError("At least two device groups are required for event calibration")
    oof_components = np.full((len(work), 4), np.nan)
    splitter = GroupKFold(n_splits=folds)
    for fold_number, (train_index, valid_index) in enumerate(
        splitter.split(base, groups=groups), start=1
    ):
        fold_estimators = _event_estimators(random_state + fold_number)
        fold_estimators, fold_imputer = _fit_event_estimators(
            fold_estimators,
            base.iloc[train_index],
            all_features.iloc[train_index],
            physics_features.iloc[train_index],
            event_target[train_index],
            None if weights is None else weights[train_index],
        )
        oof_components[valid_index] = _event_components(
            fold_estimators,
            fold_imputer,
            base.iloc[valid_index],
            all_features.iloc[valid_index],
            physics_features.iloc[valid_index],
        )

    fitted_calibrator = LogisticRegression(C=0.03, max_iter=1000)
    fitted_calibrator.fit(_logit(oof_components), event_target, sample_weight=weights)
    calibrator = FrozenLogisticCalibrator.from_estimator(fitted_calibrator)

    estimators = _event_estimators(random_state)
    estimators, imputer = _fit_event_estimators(
        estimators,
        base,
        all_features,
        physics_features,
        event_target,
        weights,
    )

    observed_near_horizon = (
        work["event_observed"].astype(bool)
        & target_rul.gt(0.0)
        & target_rul.le(float(rul_training_max_days))
    )
    rul_estimator = HistGradientBoostingRegressor(
        loss="quantile",
        quantile=quantile,
        learning_rate=0.055,
        max_iter=260,
        max_leaf_nodes=19,
        min_samples_leaf=30,
        l2_regularization=3.0,
        early_stopping=True,
        validation_fraction=0.12,
        n_iter_no_change=25,
        random_state=random_state,
    )
    rul_estimator.fit(base.loc[observed_near_horizon], target_rul.loc[observed_near_horizon])

    survivor_rul_estimator = _fit_survivor_estimator(
        base,
        target_rul,
        work["event_observed"].astype(bool),
        horizon_days,
        survivor_training_max_days,
        random_state,
    )

    balanced = None
    forecasters = None
    residual_model = None
    residual_columns = None
    trajectory_weight = 0.0
    context_estimators = None
    context_imputer = None
    context_columns = None
    hazard_estimator = None
    weights_used = (1.0, 0.0, 0.0, 0.0)
    if with_auxiliary:
        # Diverse views of the same 82 failures: reweighted, context-aware, and time-to-event.
        balanced = fit_event_time_model(
            training_snapshots,
            quantile=quantile,
            horizon_days=horizon_days,
            rul_training_max_days=rul_training_max_days,
            survivor_training_max_days=survivor_training_max_days,
            dataset_revision=dataset_revision,
            random_state=random_state,
            calibration_folds=calibration_folds,
            device_balanced=True,
            with_auxiliary=False,
        )
        context_frame = pd.concat([base, _context_features(work)], axis=1)
        context_columns = list(context_frame.columns)
        context_estimators, context_imputer = _fit_context(
            context_frame, event_target, random_state
        )
        hazard_estimator = _fit_hazard(
            base, target_rul, work["event_observed"].astype(bool), random_state
        )
        weights_used = tuple(float(value) for value in blend_weights)
        if trajectory is not None:
            forecasters, residual_model, residual_columns, trajectory_weight = _fit_trajectory_stack(
                work, trajectory, eol_times, event_target, random_state
            )

    return EventTimeModel(
        estimators,
        imputer,
        calibrator,
        rul_estimator,
        survivor_rul_estimator,
        columns,
        all_columns,
        physics_columns,
        quantile,
        horizon_days,
        dataset_revision,
        balanced,
        context_estimators,
        context_imputer,
        context_columns,
        hazard_estimator,
        weights_used,
        forecasters,
        residual_model,
        residual_columns,
        trajectory_weight,
    )