"""ShapFeatureElimination — SHAP-based recursive feature elimination with CV."""

from __future__ import annotations

import logging
import warnings
from math import floor

import numpy as np
import pandas as pd
from joblib import Parallel, delayed
from sklearn.base import BaseEstimator, clone, is_classifier
from sklearn.feature_selection import SelectorMixin
from sklearn.metrics import check_scoring
from sklearn.model_selection import check_cv
from sklearn.model_selection._search import BaseSearchCV
from sklearn.utils.validation import check_is_fitted

from recursive_pietro._compat import (
    get_early_stopping_fit_params,
    is_early_stopping_compatible,
)
from recursive_pietro._shap import aggregate_shap_importance, compute_shap_values
from recursive_pietro._validation import (
    validate_columns_to_keep,
    validate_data,
    validate_groups,
    validate_min_features,
    validate_sample_weight,
    validate_shap_variance_penalty,
    validate_step,
    validate_target,
)

logger = logging.getLogger(__name__)


def _seed_model(model, seed: int | None) -> None:
    """Set ``random_state`` on *model* in-place if it accepts one and *seed* is not None."""
    if seed is None:
        return
    try:
        params = model.get_params(deep=False)
    except AttributeError:
        return
    if "random_state" in params:
        model.set_params(random_state=seed)


class ShapFeatureElimination(SelectorMixin, BaseEstimator):
    """Backward recursive feature elimination using SHAP importance with cross-validation.

    At each round the algorithm:

    1. (Optional) Tunes hyperparameters via an sklearn SearchCV wrapper.
    2. Runs cross-validation, fitting the model on train folds and computing
       SHAP feature importance on validation folds (with optional early stopping
       for LightGBM / XGBoost / CatBoost).
    3. Removes the ``step`` lowest-importance features.

    The process repeats until ``min_features_to_select`` features remain.
    Results from every round are stored in ``report_`` so the user can choose
    the optimal feature set post-hoc.

    This class is fully sklearn-compatible: it inherits from
    ``BaseEstimator`` and ``SelectorMixin`` so it works in ``Pipeline``,
    supports ``get_params`` / ``set_params``, ``get_support``, and
    ``get_feature_names_out``.

    Parameters
    ----------
    model : estimator or sklearn SearchCV
        A sklearn-compatible classifier or regressor, or a hyperparameter
        search wrapper (``GridSearchCV``, ``RandomizedSearchCV``, etc.).
    step : int or float, default=1
        Features to remove per round.  If int, remove exactly that many.
        If float in (0, 1), remove that fraction of remaining features
        (rounded down, minimum 1).
    min_features_to_select : int, default=1
        Minimum features to keep.  The loop stops when this is reached.
    cv : int, CV splitter, or iterable, default=5
        Cross-validation strategy (passed to ``sklearn.model_selection.check_cv``).
    scoring : str or callable, default="roc_auc"
        Metric used to evaluate model performance per fold.
    n_jobs : int, default=-1
        Number of parallel jobs for CV folds (``joblib``).
    random_state : int or None, default=None
        Random seed for reproducibility.
    early_stopping_rounds : int or None, default=None
        Boosting early-stopping patience (LightGBM / XGBoost / CatBoost only).
    eval_metric : str or None, default=None
        Metric for the model's internal early-stopping evaluation.
    shap_variance_penalty_factor : float or None, default=None
        Penalise high-variance SHAP features when ranking.
        ``penalized = mean_abs_shap - std_abs_shap * factor``.
    shap_fast_mode : bool, default=False
        Use approximate SHAP values (``TreeExplainer`` only).
    sort_columns : bool, default=True
        Sort feature columns alphabetically before fitting so that results
        are invariant to the input column order.
    verbose : int, default=0
        0 = silent, 1 = progress bar, 2 = detailed logging per round.

    Attributes
    ----------
    report_ : pd.DataFrame
        Per-round results (num_features, scores, feature sets).
    selected_features_ : list[str]
        Feature names corresponding to the best validation score.
    support_ : np.ndarray of bool
        Boolean mask over the original feature set.
    ranking_ : np.ndarray of int
        Feature ranking (1 = selected, higher = eliminated earlier).
    n_features_ : int
        Number of selected features.
    column_names_ : list[str]
        Original feature names seen during ``fit``.
    """

    def __init__(
        self,
        model,
        *,
        step: int | float = 1,
        min_features_to_select: int = 1,
        cv: int | object = 5,
        scoring: str = "roc_auc",
        n_jobs: int = -1,
        random_state: int | None = None,
        early_stopping_rounds: int | None = None,
        eval_metric: str | None = None,
        shap_variance_penalty_factor: float | None = None,
        shap_fast_mode: bool = False,
        sort_columns: bool = True,
        verbose: int = 0,
    ):
        self.model = model
        self.step = step
        self.min_features_to_select = min_features_to_select
        self.cv = cv
        self.scoring = scoring
        self.n_jobs = n_jobs
        self.random_state = random_state
        self.early_stopping_rounds = early_stopping_rounds
        self.eval_metric = eval_metric
        self.shap_variance_penalty_factor = shap_variance_penalty_factor
        self.shap_fast_mode = shap_fast_mode
        self.sort_columns = sort_columns
        self.verbose = verbose

    # ------------------------------------------------------------------
    # sklearn SelectorMixin interface
    # ------------------------------------------------------------------

    def _get_support_mask(self) -> np.ndarray:
        check_is_fitted(self, "support_")
        return self.support_

    def transform(self, X):
        """Reduce X to the selected features.

        When ``sort_columns=True`` (default), columns are sorted to match
        the order seen during ``fit`` before applying the feature mask.
        """
        check_is_fitted(self, "support_")
        if self.sort_columns and isinstance(X, pd.DataFrame):
            X = X[sorted(X.columns)]
        return super().transform(X)

    # ------------------------------------------------------------------
    # fit / transform
    # ------------------------------------------------------------------

    def fit(
        self,
        X,
        y,
        *,
        groups=None,
        sample_weight=None,
        columns_to_keep: list[str] | None = None,
        column_names: list[str] | None = None,
        **shap_kwargs,
    ):
        """Run the recursive feature elimination.

        Parameters
        ----------
        X : pd.DataFrame or np.ndarray
            Feature matrix.
        y : array-like
            Target vector.
        groups : array-like or None
            Group labels for group-aware CV splitters.
        sample_weight : array-like or None
            Per-sample weights passed to the model's ``fit``.
        columns_to_keep : list[str] or None
            Columns that must never be eliminated.
        column_names : list[str] or None
            Override column names for X.
        **shap_kwargs
            Extra keyword arguments forwarded to ``shap.Explainer``.
            Supports ``check_additivity`` and ``feature_perturbation``.

        Returns
        -------
        self
        """
        # --- Validate constructor params ---
        step = validate_step(self.step)
        min_features = validate_min_features(self.min_features_to_select)
        penalty_factor = validate_shap_variance_penalty(self.shap_variance_penalty_factor)

        if self.early_stopping_rounds is not None:
            if not isinstance(self.early_stopping_rounds, int) or self.early_stopping_rounds <= 0:
                raise ValueError(
                    f"early_stopping_rounds must be a positive int, got {self.early_stopping_rounds!r}."
                )
            if self.eval_metric is None:
                warnings.warn(
                    "early_stopping_rounds is set but eval_metric is None. "
                    "Both are required for early stopping with LightGBM/XGBoost/CatBoost."
                )
            if not is_early_stopping_compatible(self.model):
                raise ValueError(
                    "early_stopping_rounds is set but model is not compatible. "
                    "Only LightGBM, XGBoost, and CatBoost support early stopping."
                )

        use_early_stopping = (
            self.early_stopping_rounds is not None
            and self.eval_metric is not None
        )

        # --- Reproducible seed sequence ---
        if self.random_state is not None:
            rng = np.random.RandomState(self.random_state)
        else:
            rng = None

        # --- Validate data ---
        X_df = validate_data(X, column_names=column_names)
        if self.sort_columns:
            X_df = X_df[sorted(X_df.columns)]
        self.column_names_ = list(X_df.columns)
        self.feature_names_in_ = np.array(self.column_names_)
        y_s = validate_target(y, X_df.index)
        sw = validate_sample_weight(sample_weight, X_df.index)
        columns_to_keep = validate_columns_to_keep(columns_to_keep, self.column_names_)

        is_search = isinstance(self.model, BaseSearchCV)
        scorer = check_scoring(self.model, scoring=self.scoring)
        cv_splitter = check_cv(self.cv, y_s, classifier=is_classifier(self.model))
        groups = validate_groups(groups, X_df.index, cv_splitter)

        # Stopping floor: at least keep columns_to_keep
        n_keep = len(columns_to_keep) if columns_to_keep else 0
        stopping = max(min_features, n_keep)

        # --- Progress bar (lazy import) ---
        progress = None
        if self.verbose == 1:
            try:
                from tqdm.auto import tqdm
                # Rough estimate of total rounds
                n_rounds = self._estimate_rounds(len(self.column_names_), step, stopping)
                progress = tqdm(total=n_rounds, desc="Feature elimination", unit="round")
            except ImportError:
                pass

        # --- Main elimination loop ---
        remaining_features = list(self.column_names_)
        results: list[dict] = []
        round_number = 0

        while len(remaining_features) > stopping:
            round_number += 1

            current_features = list(remaining_features)
            # Build the working column set (features + columns_to_keep, deduplicated, order-preserved)
            working_columns = list(dict.fromkeys(current_features + (columns_to_keep or [])))
            current_X = X_df[working_columns]

            # Optional hyperparameter tuning
            if is_search:
                search_clone = clone(self.model)
                if rng is not None:
                    _seed_model(search_clone, rng.randint(0, 2**31))
                search_clone.fit(X=current_X, y=y_s, groups=groups)
                current_model = search_clone.estimator.set_params(**search_clone.best_params_)
            else:
                current_model = clone(self.model)

            # Derive per-fold seeds for deterministic parallel execution
            splits = list(cv_splitter.split(current_X, y_s, groups))
            if rng is not None:
                fold_seeds = rng.randint(0, 2**31, size=len(splits)).tolist()
            else:
                fold_seeds = [None] * len(splits)

            # CV: compute SHAP + scores per fold
            if use_early_stopping:
                fold_results = Parallel(n_jobs=self.n_jobs)(
                    delayed(self._fold_early_stopping)(
                        current_X, y_s, current_model, train_idx, val_idx,
                        sw, scorer, fold_seed, **shap_kwargs,
                    )
                    for (train_idx, val_idx), fold_seed in zip(splits, fold_seeds)
                )
            else:
                fold_results = Parallel(n_jobs=self.n_jobs)(
                    delayed(self._fold_standard)(
                        current_X, y_s, current_model, train_idx, val_idx,
                        sw, scorer, fold_seed, **shap_kwargs,
                    )
                    for (train_idx, val_idx), fold_seed in zip(splits, fold_seeds)
                )

            # Aggregate SHAP values across folds
            # SHAP returns 2D (n_samples, n_features) for binary/regression/modern multiclass,
            # or 3D (n_classes, n_samples, n_features) for older multiclass SHAP versions.
            first_shap = fold_results[0][0]
            if first_shap.ndim == 2:
                shap_values = np.concatenate([r[0] for r in fold_results], axis=0)
            else:  # 3D: concat along samples axis (axis=1)
                shap_values = np.concatenate([r[0] for r in fold_results], axis=1)

            scores_train = [r[1] for r in fold_results]
            scores_val = [r[2] for r in fold_results]

            importance_df = aggregate_shap_importance(
                shap_values, working_columns, penalty_factor=penalty_factor,
            )

            # Determine features to remove
            # When columns_to_keep is active, the removable set should be able
            # to go to zero (the stopping criteria already accounts for kept columns).
            effective_min = 0 if columns_to_keep else min_features
            features_to_remove = self._pick_features_to_remove(
                importance_df, step, effective_min, columns_to_keep,
            )
            remove_set = set(features_to_remove)
            remaining_features = [f for f in current_features if f not in remove_set]

            # Record round
            results.append({
                "num_features": len(current_features),
                "features_set": current_features,
                "eliminated_features": features_to_remove,
                "train_metric_mean": float(np.mean(scores_train)),
                "train_metric_std": float(np.std(scores_train)),
                "val_metric_mean": float(np.mean(scores_val)),
                "val_metric_std": float(np.std(scores_val)),
            })

            if self.verbose >= 2:
                logger.info(
                    "Round %d: %d features, val=%.4f +/- %.4f, removed %s",
                    round_number, len(current_features),
                    results[-1]["val_metric_mean"], results[-1]["val_metric_std"],
                    features_to_remove,
                )
            if progress is not None:
                progress.update(1)

        if progress is not None:
            progress.close()

        # --- Build report ---
        self.report_ = pd.DataFrame(results, index=range(1, len(results) + 1))

        # --- Derive sklearn-compatible attributes ---
        self._set_selection_attributes()

        return self

    # ------------------------------------------------------------------
    # Feature set selection
    # ------------------------------------------------------------------

    def get_feature_set(
        self,
        method: str | int = "best",
        threshold: float = 1.0,
    ) -> list[str]:
        """Return a feature set from the elimination report.

        Parameters
        ----------
        method : str or int
            - ``"best"``: highest validation score.
            - ``"best_coherent"``: within ``threshold`` of best score,
              pick the round with lowest validation std.
            - ``"best_parsimonious"``: within ``threshold`` of best score,
              pick the round with fewest features.
            - int: return the feature set from the round with exactly
              that many features.
        threshold : float
            Score margin for coherent / parsimonious methods.

        Returns
        -------
        list[str]
        """
        check_is_fitted(self, "report_")
        report = self.report_

        if isinstance(method, (int, np.integer)):
            method = int(method)
            match = report[report["num_features"] == method]
            if match.empty:
                valid = sorted(report["num_features"].unique())
                raise ValueError(
                    f"No round had exactly {method} features. Available: {valid}"
                )
            return match.iloc[0]["features_set"]

        if method == "best":
            idx = report["val_metric_mean"].idxmax()
        elif method == "best_coherent":
            best_score = report["val_metric_mean"].max()
            within = report[report["val_metric_mean"] >= best_score - threshold]
            idx = within["val_metric_std"].idxmin()
        elif method == "best_parsimonious":
            best_score = report["val_metric_mean"].max()
            within = report[report["val_metric_mean"] >= best_score - threshold]
            idx = within["num_features"].idxmin()
        else:
            raise ValueError(
                f"method must be 'best', 'best_coherent', 'best_parsimonious', or int; got {method!r}"
            )

        return report.loc[idx, "features_set"]

    # ------------------------------------------------------------------
    # Plotting
    # ------------------------------------------------------------------

    def plot(self, *, show: bool = True, **figure_kwargs):
        """Plot train/validation performance across elimination rounds.

        Requires matplotlib (optional dependency).

        Returns
        -------
        matplotlib.figure.Figure
        """
        check_is_fitted(self, "report_")
        try:
            import matplotlib.pyplot as plt
        except ImportError as e:
            raise ImportError(
                "matplotlib is required for plotting. Install it with: "
                "pip install recursive-pietro[plot]"
            ) from e

        report = self.report_
        n_feat = report["num_features"]
        fig, ax = plt.subplots(**figure_kwargs)

        ax.plot(n_feat, report["train_metric_mean"], label="Train")
        ax.fill_between(
            n_feat,
            report["train_metric_mean"] - report["train_metric_std"],
            report["train_metric_mean"] + report["train_metric_std"],
            alpha=0.3,
        )
        ax.plot(n_feat, report["val_metric_mean"], label="Validation")
        ax.fill_between(
            n_feat,
            report["val_metric_mean"] - report["val_metric_std"],
            report["val_metric_mean"] + report["val_metric_std"],
            alpha=0.3,
        )

        ax.set_xlabel("Number of features")
        ax.set_ylabel(f"Score ({self.scoring})")
        ax.set_title("SHAP Recursive Feature Elimination")
        ax.legend(loc="lower left")
        ax.invert_xaxis()
        ax.set_xticks(list(reversed(n_feat.tolist())))

        if show:
            plt.show()
        else:
            plt.close(fig)

        return fig

    # ------------------------------------------------------------------
    # Private helpers
    # ------------------------------------------------------------------

    def _fold_standard(
        self, X, y, model, train_idx, val_idx, sample_weight, scorer, fold_seed, **shap_kwargs,
    ):
        """Fit model, score, compute SHAP — standard (no early stopping) path."""
        model = clone(model)
        _seed_model(model, fold_seed)
        X_train, X_val = X.iloc[train_idx], X.iloc[val_idx]
        y_train, y_val = y.iloc[train_idx], y.iloc[val_idx]

        fit_kwargs = {}
        if sample_weight is not None:
            fit_kwargs["sample_weight"] = sample_weight.iloc[train_idx]

        model.fit(X_train, y_train, **fit_kwargs)

        score_train = scorer(model, X_train, y_train)
        score_val = scorer(model, X_val, y_val)

        shap_vals = compute_shap_values(
            model, X_val,
            approximate=self.shap_fast_mode,
            random_state=fold_seed,
            verbose=self.verbose,
            **shap_kwargs,
        )
        return shap_vals, score_train, score_val

    def _fold_early_stopping(
        self, X, y, model, train_idx, val_idx, sample_weight, scorer, fold_seed, **shap_kwargs,
    ):
        """Fit model with early stopping, score, compute SHAP."""
        model = clone(model)
        _seed_model(model, fold_seed)
        X_train, X_val = X.iloc[train_idx], X.iloc[val_idx]
        y_train, y_val = y.iloc[train_idx], y.iloc[val_idx]

        fit_params = get_early_stopping_fit_params(
            model, X_train, y_train, X_val, y_val,
            early_stopping_rounds=self.early_stopping_rounds,
            eval_metric=self.eval_metric,
            sample_weight=sample_weight,
            train_index=train_idx,
            val_index=val_idx,
            verbose=self.verbose,
        )
        model.fit(**fit_params)

        score_train = scorer(model, X_train, y_train)
        score_val = scorer(model, X_val, y_val)

        shap_vals = compute_shap_values(
            model, X_val,
            approximate=self.shap_fast_mode,
            random_state=fold_seed,
            verbose=self.verbose,
            **shap_kwargs,
        )
        return shap_vals, score_train, score_val

    @staticmethod
    def _pick_features_to_remove(
        importance_df: pd.DataFrame,
        step: int | float,
        min_features: int,
        columns_to_keep: list[str] | None,
    ) -> list[str]:
        """Select the least important features to remove this round."""
        # Exclude protected columns from candidates
        if columns_to_keep:
            mask = importance_df.index.isin(columns_to_keep)
            candidates = importance_df[~mask]
        else:
            candidates = importance_df

        n_candidates = len(candidates)

        if isinstance(step, int):
            n_remove = step
        else:
            n_remove = max(1, int(floor(n_candidates * step)))

        # Don't drop below min_features (among removable features)
        max_removable = n_candidates - min_features
        n_remove = max(0, min(n_remove, max_removable))

        if n_remove == 0:
            return []

        # Bottom n_remove features (lowest importance) — importance_df is sorted descending
        return candidates.iloc[-n_remove:].index.tolist()

    @staticmethod
    def _estimate_rounds(n_features: int, step: int | float, stopping: int) -> int:
        """Rough estimate of total elimination rounds (for progress bar)."""
        n = n_features
        rounds = 0
        while n > stopping:
            rounds += 1
            if isinstance(step, int):
                n -= max(1, min(step, n - stopping))
            else:
                remove = max(1, int(floor(n * step)))
                n -= max(1, min(remove, n - stopping))
        return max(rounds, 1)

    def _set_selection_attributes(self):
        """Compute support_, ranking_, selected_features_, n_features_ from report_."""
        report = self.report_
        best_idx = report["val_metric_mean"].idxmax()
        self.selected_features_ = report.loc[best_idx, "features_set"]
        self.n_features_ = len(self.selected_features_)

        selected_set = set(self.selected_features_)
        self.support_ = np.array([c in selected_set for c in self.column_names_])

        # Build ranking: 1 = selected, higher = eliminated earlier in time.
        # Walk elimination rounds in order. Features eliminated before the best
        # round are NOT selected → they get a rank based on elimination order.
        # Features in the selected set always get rank 1.
        best_num = report.loc[best_idx, "num_features"]
        eliminated_before_best = []
        for _, row in report.iterrows():
            if row["num_features"] > best_num:
                eliminated_before_best.extend(row["eliminated_features"])

        ranking = {}
        for feat in selected_set:
            ranking[feat] = 1
        # Eliminated features: first eliminated = highest rank
        # (eliminated_before_best is in chronological order, first eliminated = earliest)
        for i, feat in enumerate(reversed(eliminated_before_best)):
            ranking[feat] = 2 + i

        # Any remaining features not yet ranked (shouldn't happen, but defensive)
        for feat in self.column_names_:
            if feat not in ranking:
                ranking[feat] = 1

        self.ranking_ = np.array([ranking[c] for c in self.column_names_])
