ml/feature_selection.py

"""
ml/feature_selection.py — SHAP-based feature selection & importance analysis.

Capabilities:
  1. SHAP value computation for tree-based models (XGB/LGB)
  2. Drop redundant and noisy features
  3. Regime-specific feature importance analysis
  4. Stable feature set extraction across walk-forward windows
"""
from __future__ import annotations

import logging
from dataclasses import dataclass, field

import numpy as np
import pandas as pd

logger = logging.getLogger("trading_system.ml.feature_selection")


@dataclass
class FeatureImportance:
    """Results of feature importance analysis."""
    feature_names: list[str]
    mean_abs_shap: np.ndarray            # mean |SHAP| per feature
    std_shap: np.ndarray                 # std of SHAP per feature
    regime_importance: dict[str, np.ndarray] = field(default_factory=dict)
    selected_features: list[str] = field(default_factory=list)
    dropped_features: list[str] = field(default_factory=list)


class SHAPFeatureSelector:
    """SHAP-based feature selection with stability analysis.

    Strategy:
      - Compute SHAP values for fitted tree model
      - Rank features by mean |SHAP|
      - Drop features with negligible contribution
      - Remove highly correlated redundant features
      - Optionally analyze per-regime importance
    """

    def __init__(
        self,
        min_importance_pct: float = 0.5,    # drop features < 0.5% of max importance
        max_correlation: float = 0.90,       # drop correlated pairs
        min_features: int = 10,              # keep at least this many
        max_features: int = 50,              # cap features
        stability_threshold: float = 0.3,    # feature must appear in 30% of WF windows
    ):
        self.min_importance_pct = min_importance_pct
        self.max_correlation = max_correlation
        self.min_features = min_features
        self.max_features = max_features
        self.stability_threshold = stability_threshold
        self._window_selections: list[list[str]] = []

    def compute_importance(
        self,
        model,
        X: pd.DataFrame,
        n_samples: int = 500,
    ) -> FeatureImportance:
        """Compute SHAP-based feature importance.

        Args:
            model: Fitted tree model (XGB/LGB/sklearn).
            X: Feature DataFrame.
            n_samples: Max samples for SHAP computation (speed).

        Returns:
            FeatureImportance with rankings.
        """
        import shap

        feature_names = list(X.columns)

        # Subsample for speed
        if len(X) > n_samples:
            X_sample = X.sample(n=n_samples, random_state=42)
        else:
            X_sample = X

        try:
            explainer = shap.TreeExplainer(model)
            shap_values = explainer.shap_values(X_sample)
        except Exception:
            # Fallback to permutation explainer
            try:
                explainer = shap.Explainer(model, X_sample)
                shap_obj = explainer(X_sample)
                shap_values = shap_obj.values
            except Exception as e:
                logger.warning("SHAP computation failed: %s. Using model importances.", e)
                return self._fallback_importance(model, feature_names)

        # Handle multi-class SHAP (list of arrays)
        if isinstance(shap_values, list):
            # Average across classes
            stacked = np.stack([np.abs(sv) for sv in shap_values], axis=0)
            abs_shap = stacked.mean(axis=0)  # (n_samples, n_features)
        elif shap_values.ndim == 3:
            abs_shap = np.abs(shap_values).mean(axis=2)
        else:
            abs_shap = np.abs(shap_values)

        mean_abs = abs_shap.mean(axis=0)
        std_vals = abs_shap.std(axis=0)

        return FeatureImportance(
            feature_names=feature_names,
            mean_abs_shap=mean_abs,
            std_shap=std_vals,
        )

    def _fallback_importance(self, model, feature_names) -> FeatureImportance:
        """Use model's built-in feature importance if SHAP fails."""
        try:
            imp = model.feature_importances_
        except AttributeError:
            imp = np.ones(len(feature_names))

        return FeatureImportance(
            feature_names=feature_names,
            mean_abs_shap=imp / imp.sum() if imp.sum() > 0 else imp,
            std_shap=np.zeros(len(feature_names)),
        )

    def select_features(
        self,
        importance: FeatureImportance,
        X: pd.DataFrame,
    ) -> FeatureImportance:
        """Select features based on importance and correlation.

        Returns updated FeatureImportance with selected/dropped lists.
        """
        feat_names = importance.feature_names
        scores = importance.mean_abs_shap

        # Rank by importance
        sorted_idx = np.argsort(-scores)
        threshold = scores.max() * (self.min_importance_pct / 100.0) if scores.max() > 0 else 0

        # Keep features above threshold
        kept_idx = [i for i in sorted_idx if scores[i] >= threshold]

        # Enforce min/max
        if len(kept_idx) < self.min_features:
            kept_idx = list(sorted_idx[:self.min_features])
        if len(kept_idx) > self.max_features:
            kept_idx = kept_idx[:self.max_features]

        kept_names = [feat_names[i] for i in kept_idx]

        # Remove correlated features (keep the more important one)
        kept_names = self._remove_correlated(X, kept_names, scores, feat_names)

        if len(kept_names) < self.min_features:
            # Re-add top features if too many removed
            for i in sorted_idx:
                if feat_names[i] not in kept_names:
                    kept_names.append(feat_names[i])
                if len(kept_names) >= self.min_features:
                    break

        dropped = [f for f in feat_names if f not in kept_names]

        importance.selected_features = kept_names
        importance.dropped_features = dropped

        # Record for stability tracking
        self._window_selections.append(kept_names)

        logger.info(
            "Selected %d/%d features (dropped %d)",
            len(kept_names), len(feat_names), len(dropped),
        )
        return importance

    def _remove_correlated(
        self,
        X: pd.DataFrame,
        kept_names: list[str],
        scores: np.ndarray,
        all_names: list[str],
    ) -> list[str]:
        """Remove highly correlated features, keeping the more important one."""
        if len(kept_names) < 2:
            return kept_names

        subset = X[kept_names].copy()
        corr = subset.corr().abs()

        to_drop = set()
        for i in range(len(kept_names)):
            if kept_names[i] in to_drop:
                continue
            for j in range(i + 1, len(kept_names)):
                if kept_names[j] in to_drop:
                    continue
                if corr.iloc[i, j] > self.max_correlation:
                    # Drop the less important one
                    idx_i = all_names.index(kept_names[i])
                    idx_j = all_names.index(kept_names[j])
                    if scores[idx_i] >= scores[idx_j]:
                        to_drop.add(kept_names[j])
                    else:
                        to_drop.add(kept_names[i])

        return [f for f in kept_names if f not in to_drop]

    def get_stable_features(self) -> list[str]:
        """Get features that appear consistently across walk-forward windows."""
        if not self._window_selections:
            return []

        n_windows = len(self._window_selections)
        feature_counts: dict[str, int] = {}
        for selection in self._window_selections:
            for feat in selection:
                feature_counts[feat] = feature_counts.get(feat, 0) + 1

        threshold = max(1, int(n_windows * self.stability_threshold))
        stable = [f for f, c in feature_counts.items() if c >= threshold]

        # Sort by frequency
        stable.sort(key=lambda f: -feature_counts.get(f, 0))
        return stable

    def regime_importance(
        self,
        model,
        X: pd.DataFrame,
        regime_labels: np.ndarray,
        regime_names: dict[int, str] | None = None,
    ) -> dict[str, np.ndarray]:
        """Compute SHAP importance per regime.

        Args:
            model: Fitted tree model.
            X: Feature DataFrame.
            regime_labels: Array of regime integers per row.
            regime_names: Optional mapping int→name.

        Returns:
            Dict of regime_name → mean |SHAP| array (per feature).
        """
        import shap

        if regime_names is None:
            regime_names = {0: "SIDEWAYS", 1: "BULL", 2: "BEAR"}

        result = {}
        for r_int, r_name in regime_names.items():
            mask = regime_labels == r_int
            if mask.sum() < 50:
                continue

            X_regime = X[mask]
            if len(X_regime) > 300:
                X_regime = X_regime.sample(n=300, random_state=42)

            try:
                explainer = shap.TreeExplainer(model)
                sv = explainer.shap_values(X_regime)
                if isinstance(sv, list):
                    abs_sv = np.mean([np.abs(s) for s in sv], axis=0)
                else:
                    abs_sv = np.abs(sv)
                result[r_name] = abs_sv.mean(axis=0)
            except Exception as e:
                logger.warning("Regime SHAP for %s failed: %s", r_name, e)

        return result


def print_feature_report(importance: FeatureImportance, top_n: int = 20):
    """Print a feature importance summary."""
    names = importance.feature_names
    scores = importance.mean_abs_shap

    sorted_idx = np.argsort(-scores)

    print("\n" + "=" * 60)
    print("  SHAP FEATURE IMPORTANCE REPORT")
    print("=" * 60)

    for rank, i in enumerate(sorted_idx[:top_n], 1):
        bar = "█" * int(scores[i] / (scores.max() + 1e-9) * 30)
        selected = "✓" if names[i] in importance.selected_features else "✗"
        print(f"  {rank:2d}. [{selected}] {names[i]:<30s}  {scores[i]:.4f}  {bar}")

    if importance.selected_features:
        print(f"\n  Selected: {len(importance.selected_features)} features")
        print(f"  Dropped:  {len(importance.dropped_features)} features")

    if importance.regime_importance:
        print("\n  --- Per-Regime Top-5 ---")
        for regime, imp in importance.regime_importance.items():
            top5 = np.argsort(-imp)[:5]
            feats = ", ".join(names[i] for i in top5)
            print(f"  {regime}: {feats}")
    print("=" * 60)