src/models/ensemble.py

"""
Stacking Ensemble for fraud detection.

Combines predictions from multiple base models (XGBoost, LightGBM, deep learning)
using a meta-learner (logistic regression) trained on out-of-fold predictions
to avoid data leakage and overfitting.

Supports two stacking strategies:
  - "stack" (default): True out-of-fold stacking. Base models are retrained
    per fold so OOF predictions never see their own training data.
  - "blend": Legacy mode. Pre-trained base models generate predictions
    directly (no per-fold retraining). Faster but prone to leakage.
"""

import copy
import logging
import time
from typing import Any, Callable, Dict, List, Optional, Tuple, Union

import numpy as np
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import average_precision_score, roc_auc_score
from sklearn.model_selection import StratifiedKFold

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Type alias for a model factory: () -> model_instance
# ---------------------------------------------------------------------------
ModelFactory = Callable[[], Any]


class StackingEnsemble:
    """Stacking ensemble that combines multiple fraud detection models.

    Uses cross-validated out-of-fold predictions as meta-features to train a
    second-level logistic regression model.

    Two construction patterns are supported:

    1. **True stacking (recommended)** -- pass *model_factories* so the
       ensemble can create fresh model instances for each CV fold::

           ensemble = StackingEnsemble(
               model_factories={
                   "xgb": lambda: FraudXGBoost(xgb_config),
                   "lgb": lambda: FraudLightGBM(lgb_config),
               },
           )
           ensemble.train(X_train, y_train, X_val, y_val, method="stack")

    2. **Blend mode (legacy)** -- pass pre-trained *base_models* and call
       ``train(..., method="blend")``::

           ensemble = StackingEnsemble(base_models={"xgb": trained_xgb, ...})
           ensemble.train(X_train, y_train, X_val, y_val, method="blend")

    Both patterns expose the same ``predict_proba`` / ``predict`` interface
    after training.
    """

    def __init__(
        self,
        base_models: Optional[Dict[str, Any]] = None,
        model_factories: Optional[Dict[str, ModelFactory]] = None,
        config: Optional[dict] = None,
    ):
        """Initialize the stacking ensemble.

        Args:
            base_models: Dict mapping model name to a *pre-trained* model
                instance.  Each must implement ``predict_proba(X)``.  Used in
                blend mode or as a fallback when *model_factories* is not
                provided.
            model_factories: Dict mapping model name to a callable that returns
                a **new, untrained** model instance.  Each returned instance
                must implement ``train(X_train, y_train, X_val, y_val)`` and
                ``predict_proba(X)``.  Required for true stacking mode.
            config: Optional ensemble configuration.  Recognised keys:

                * ``cv_folds`` (int, default 5): number of stratified folds.
                * ``meta_C`` (float, default 1.0): regularisation strength
                  for the logistic-regression meta-learner.
                * ``random_state`` (int, default 42): random seed.
        """
        if base_models is None and model_factories is None:
            raise ValueError(
                "At least one of 'base_models' or 'model_factories' must be provided."
            )

        self.base_models: Dict[str, Any] = base_models or {}
        self.model_factories: Dict[str, ModelFactory] = model_factories or {}
        self.config = config or {}

        self.meta_learner: Optional[LogisticRegression] = None
        self.cv_folds: int = self.config.get("cv_folds", 5)
        self._random_state: int = self.config.get("random_state", 42)
        self._meta_C: float = self.config.get("meta_C", 1.0)

        # Populated after train(); holds the final base models used at
        # inference time (retrained on full training set for stacking mode,
        # or the original pre-trained models for blend mode).
        self._inference_models: Dict[str, Any] = {}

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def train(
        self,
        X_train: np.ndarray,
        y_train: np.ndarray,
        X_val: np.ndarray,
        y_val: np.ndarray,
        method: str = "stack",
        feature_names: Optional[list] = None,
    ) -> None:
        """Train the stacking ensemble.

        Args:
            X_train: Training feature matrix.
            y_train: Training labels (binary).
            X_val: Validation feature matrix (used for early-stopping inside
                base models and for diagnostic evaluation of the ensemble).
            y_val: Validation labels.
            method: ``"stack"`` for true out-of-fold stacking (default),
                ``"blend"`` for legacy pre-trained-model mode.
            feature_names: Optional feature column names forwarded to base
                model ``train()`` calls.  Ignored in blend mode.
        """
        method = method.lower()
        if method not in ("stack", "blend"):
            raise ValueError(f"Unknown method '{method}'. Choose 'stack' or 'blend'.")

        if method == "stack":
            if not self.model_factories:
                raise ValueError(
                    "True stacking (method='stack') requires 'model_factories'. "
                    "Pass pre-trained 'base_models' and use method='blend', or "
                    "provide model_factories."
                )
            oof_predictions = self._train_stack(
                X_train, y_train, X_val, y_val, feature_names,
            )
        else:
            if not self.base_models:
                raise ValueError(
                    "Blend mode requires pre-trained 'base_models'."
                )
            oof_predictions = self._train_blend(X_train, y_train)

        # ----- Fit the meta-learner on OOF predictions -----
        logger.info("Fitting meta-learner on OOF predictions of shape %s", oof_predictions.shape)
        self.meta_learner = LogisticRegression(
            C=self._meta_C,
            max_iter=1000,
            class_weight="balanced",
            random_state=self._random_state,
        )
        self.meta_learner.fit(oof_predictions, y_train)

        # Log meta-learner coefficients (model importance)
        model_names = list(self._inference_models.keys())
        coefs = self.meta_learner.coef_[0]
        for name, coef in zip(model_names, coefs):
            logger.info("  Meta-learner weight for %s: %.4f", name, coef)

        # ----- Evaluate on held-out validation set -----
        val_meta_features = self._get_meta_features(X_val)
        val_proba = self.meta_learner.predict_proba(val_meta_features)[:, 1]
        val_auc = roc_auc_score(y_val, val_proba)
        val_ap = average_precision_score(y_val, val_proba)
        logger.info(
            "Ensemble validation -- AUC-ROC: %.4f, Avg Precision: %.4f",
            val_auc,
            val_ap,
        )

    def predict_proba(self, X: np.ndarray) -> np.ndarray:
        """Predict fraud probability using the stacking ensemble.

        Args:
            X: Feature matrix.

        Returns:
            1-D array of fraud probability scores.

        Raises:
            RuntimeError: If the ensemble has not been trained yet.
        """
        if self.meta_learner is None:
            raise RuntimeError("Ensemble has not been trained. Call train() first.")
        meta_features = self._get_meta_features(X)
        return self.meta_learner.predict_proba(meta_features)[:, 1]

    def predict(self, X: np.ndarray, threshold: float = 0.5) -> np.ndarray:
        """Predict binary fraud labels.

        Args:
            X: Feature matrix.
            threshold: Decision threshold for the positive class.

        Returns:
            Binary prediction array (0 or 1).
        """
        return (self.predict_proba(X) >= threshold).astype(int)

    # ------------------------------------------------------------------
    # Internal: true stacking
    # ------------------------------------------------------------------

    def _train_stack(
        self,
        X_train: np.ndarray,
        y_train: np.ndarray,
        X_val: np.ndarray,
        y_val: np.ndarray,
        feature_names: Optional[list],
    ) -> np.ndarray:
        """Run K-fold stacking: retrain base models per fold.

        Returns the (n_samples, n_models) OOF prediction matrix.
        After the OOF loop, retrains each base model on the full training set
        so that ``_inference_models`` are ready for ``predict_proba``.
        """
        model_names = list(self.model_factories.keys())
        n_models = len(model_names)
        n_samples = len(X_train)

        logger.info(
            "Training stacking ensemble (method=stack) with %d base models, "
            "%d CV folds, %d training samples",
            n_models,
            self.cv_folds,
            n_samples,
        )

        oof_predictions = np.zeros((n_samples, n_models))
        kf = StratifiedKFold(
            n_splits=self.cv_folds, shuffle=True, random_state=self._random_state,
        )

        for fold_idx, (train_idx, val_idx) in enumerate(kf.split(X_train, y_train)):
            fold_start = time.perf_counter()
            logger.info(
                "  Fold %d/%d  --  train: %d samples, val: %d samples",
                fold_idx + 1,
                self.cv_folds,
                len(train_idx),
                len(val_idx),
            )

            X_fold_train, y_fold_train = X_train[train_idx], y_train[train_idx]
            X_fold_val, y_fold_val = X_train[val_idx], y_train[val_idx]

            for model_idx, name in enumerate(model_names):
                model = self.model_factories[name]()
                logger.info("    Training '%s' on fold %d ...", name, fold_idx + 1)
                model.train(
                    X_fold_train,
                    y_fold_train,
                    X_fold_val,
                    y_fold_val,
                    feature_names=feature_names,
                )
                oof_predictions[val_idx, model_idx] = model.predict_proba(X_fold_val)

            fold_elapsed = time.perf_counter() - fold_start
            logger.info("  Fold %d complete in %.1fs", fold_idx + 1, fold_elapsed)

        # Retrain each base model on the FULL training set for inference
        logger.info("Retraining all base models on the full training set for inference ...")
        self._inference_models = {}
        for name in model_names:
            model = self.model_factories[name]()
            logger.info("  Retraining '%s' on %d samples ...", name, n_samples)
            model.train(
                X_train,
                y_train,
                X_val,
                y_val,
                feature_names=feature_names,
            )
            self._inference_models[name] = model

        return oof_predictions

    # ------------------------------------------------------------------
    # Internal: blend (legacy) mode
    # ------------------------------------------------------------------

    def _train_blend(
        self,
        X_train: np.ndarray,
        y_train: np.ndarray,
    ) -> np.ndarray:
        """Generate OOF-style predictions from pre-trained models (blend).

        In blend mode the base models are already trained on the full training
        set, so the OOF predictions are *not* truly out-of-fold.  This is kept
        for backward compatibility; prefer ``method='stack'`` for proper
        stacking.

        Returns the (n_samples, n_models) prediction matrix.
        """
        model_names = list(self.base_models.keys())
        n_models = len(model_names)

        logger.info(
            "Training stacking ensemble (method=blend) with %d pre-trained base models, "
            "%d CV folds",
            n_models,
            self.cv_folds,
        )

        oof_predictions = np.zeros((len(X_train), n_models))
        kf = StratifiedKFold(
            n_splits=self.cv_folds, shuffle=True, random_state=self._random_state,
        )

        for fold_idx, (train_idx, val_idx) in enumerate(kf.split(X_train, y_train)):
            logger.info("  Fold %d/%d", fold_idx + 1, self.cv_folds)
            X_fold_val = X_train[val_idx]
            for model_idx, (name, model) in enumerate(self.base_models.items()):
                oof_predictions[val_idx, model_idx] = model.predict_proba(X_fold_val)

        # In blend mode the pre-trained models are used directly at inference.
        self._inference_models = dict(self.base_models)

        return oof_predictions

    # ------------------------------------------------------------------
    # Shared helpers
    # ------------------------------------------------------------------

    def _get_meta_features(self, X: np.ndarray) -> np.ndarray:
        """Generate meta-features from the inference base-model predictions.

        Args:
            X: Raw feature matrix.

        Returns:
            (n_samples, n_models) array of base-model probability predictions.
        """
        if not self._inference_models:
            raise RuntimeError(
                "No inference models available. Call train() first."
            )
        meta_features = np.column_stack(
            [model.predict_proba(X) for model in self._inference_models.values()]
        )
        return meta_features