src/features/selector.py

"""Feature selection, implemented so that fitting outside a CV fold is impossible.

THE CONSTRAINT. n is 124 cells. Selecting features on the full dataset is the
fastest way to manufacture a result that looks strong and does not replicate:
the selector sees the labels of the cells it will later be evaluated on, and
"which features correlate with the target" is exactly the information a test
fold is supposed to withhold.

HOW IT IS ENFORCED HERE. Every selector below is an sklearn transformer with a
``fit``/``transform`` split. Composed into a ``Pipeline`` and passed to
``cross_val_score`` or ``GridSearchCV``, ``fit`` is called on the training fold
ONLY -- scikit-learn guarantees that, so the guarantee does not depend on this
project remembering to do it. There is deliberately no module-level
"select_features(X, y)" convenience function, because that is the shape of API
that invites full-data use.

``configs/features.yaml`` sets ``selection.fit_inside_cv_only: true`` and the
schema in ``src/utils/config.py`` REJECTS the config if it is ever false.
"""

from __future__ import annotations

import numpy as np
import pandas as pd
from scipy import stats
from sklearn.base import BaseEstimator, TransformerMixin
from sklearn.feature_selection import mutual_info_classif, mutual_info_regression
from sklearn.impute import SimpleImputer
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler
from sklearn.utils.validation import check_is_fitted

from src.utils.config import ConfigBundle


def _abs_spearman(x: np.ndarray, y: np.ndarray) -> float:
    """|Spearman rho| between two aligned 1-D arrays, 0.0 if undefined.

    Takes arrays rather than Series on purpose: pandas correlates on the shared
    INDEX, which silently mis-pairs a bootstrap subset (original row labels)
    against a freshly built target (RangeIndex).
    """
    if x.size != y.size or x.size < 3:
        return 0.0
    finite = np.isfinite(x) & np.isfinite(y)
    if finite.sum() < 3:
        return 0.0
    x, y = x[finite], y[finite]
    if np.ptp(x) == 0 or np.ptp(y) == 0:
        return 0.0
    rho = stats.spearmanr(x, y).statistic
    return float(abs(rho)) if np.isfinite(rho) else 0.0


def _encode_target(y) -> tuple[np.ndarray, bool]:
    """Return a numeric target and whether it is categorical.

    The selectors serve both routes of the project: regression on log10 cycle
    life, and the secondary direct classification on grade labels. Grades arrive
    as strings, which neither mutual information nor a rank correlation can
    consume.

    Sorted order is used for the encoding, which for the grade labels A/B/C
    yields 0/1/2 -- their true ordinal order, most demanding first. That is
    meaningful rather than arbitrary here, so a rank statistic computed on the
    codes measures something real.
    """
    array = np.asarray(y).ravel()
    if array.dtype.kind in "biufc":
        return array.astype("float64"), False
    categories = np.unique(array)
    lookup = {value: index for index, value in enumerate(categories)}
    return np.array([lookup[v] for v in array], dtype="float64"), True


def _abs_spearman_matrix(X: np.ndarray, y: np.ndarray) -> np.ndarray:
    """|Spearman rho| between every column of ``X`` and ``y``, vectorised.

    Spearman is Pearson on ranks, so ranking each column and then taking a
    single centred, normalised matrix product is EXACTLY equivalent to calling
    a Spearman routine per column -- it is a speed change, not an approximation.
    Ties are handled by ``rankdata``'s average method, matching
    ``scipy.stats.spearmanr``.

    Columns that are constant, or all-NaN, score 0.0 rather than NaN: they carry
    no information, and a NaN would propagate into the ranking.
    """
    n_rows, n_cols = X.shape
    if n_rows < 3:
        return np.zeros(n_cols)

    # NaNs cannot be ranked. Column medians keep the column usable without
    # inventing signal, and a fully-NaN column is zeroed out below anyway.
    if np.isnan(X).any():
        medians = np.nanmedian(np.where(np.isfinite(X), X, np.nan), axis=0)
        medians = np.where(np.isfinite(medians), medians, 0.0)
        X = np.where(np.isfinite(X), X, medians)

    ranks_x = stats.rankdata(X, axis=0)
    ranks_y = stats.rankdata(y)

    xc = ranks_x - ranks_x.mean(axis=0)
    yc = ranks_y - ranks_y.mean()

    x_norm = np.sqrt((xc**2).sum(axis=0))
    y_norm = np.sqrt((yc**2).sum())

    with np.errstate(invalid="ignore", divide="ignore"):
        rho = (xc * yc[:, None]).sum(axis=0) / (x_norm * y_norm)
    return np.nan_to_num(np.abs(rho), nan=0.0, posinf=0.0, neginf=0.0)


class _FrameAwareSelector(BaseEstimator, TransformerMixin):
    """Base class that preserves feature names across transforms."""

    def _record_names(self, X) -> np.ndarray:
        if isinstance(X, pd.DataFrame):
            self.feature_names_in_ = np.asarray(X.columns)
        else:
            self.feature_names_in_ = np.asarray([f"x{i}" for i in range(np.shape(X)[1])])
        return self.feature_names_in_

    def _subset(self, X):
        if isinstance(X, pd.DataFrame):
            return X.loc[:, self.support_]
        return np.asarray(X)[:, self.support_]

    def get_feature_names_out(self, input_features=None) -> np.ndarray:
        check_is_fitted(self, "support_")
        return self.feature_names_in_[self.support_]


class VarianceFilter(_FrameAwareSelector):
    """Drop near-constant and all-missing columns.

    A column with no variance in the TRAINING fold cannot inform a model fitted
    on that fold. This also removes the degenerate columns the builder
    deliberately leaves in place (``policy_c3`` is all-NaN in this corpus and
    ``policy_n_steps`` is constant), which is where that decision belongs.
    """

    def __init__(self, threshold: float = 1e-8):
        self.threshold = threshold

    def fit(self, X, y=None):
        names = self._record_names(X)
        frame = pd.DataFrame(X, columns=names)
        variances = frame.var(axis=0, skipna=True)
        all_missing = frame.isna().all(axis=0)
        self.support_ = (variances > self.threshold).to_numpy() & ~all_missing.to_numpy()
        if not self.support_.any():  # never hand an empty matrix downstream
            self.support_ = np.ones(len(names), dtype=bool)
        return self

    def transform(self, X):
        check_is_fitted(self, "support_")
        return self._subset(X)


class CorrelationFilter(_FrameAwareSelector):
    """Drop one of every pair of features correlated above a threshold.

    WHY it matters here specifically: the curve features are collinear by
    construction -- ``dq_var`` and ``log_abs_dq_var`` are monotone transforms of
    one another, and several ΔQ(V) statistics move together. Severe collinearity
    destabilises linear-model coefficients, which in turn makes the SHAP
    attributions of Phase 9 unstable, and an unstable attribution cannot support
    an auditable scrap decision.

    Of each correlated pair the LATER column is dropped, so the ordering of the
    feature registry determines precedence deterministically rather than by
    whichever the correlation matrix happened to list first.
    """

    def __init__(self, threshold: float = 0.95):
        self.threshold = threshold

    def fit(self, X, y=None):
        names = self._record_names(X)
        frame = pd.DataFrame(X, columns=names)
        # .copy() because pandas returns a read-only view here, which
        # np.fill_diagonal cannot write into.
        corr = frame.corr(method="spearman").abs().to_numpy().copy()
        np.fill_diagonal(corr, 0.0)
        corr = np.nan_to_num(corr, nan=0.0)

        keep = np.ones(len(names), dtype=bool)
        for j in range(len(names)):
            if not keep[j]:
                continue
            for k in range(j + 1, len(names)):
                if keep[k] and corr[j, k] > self.threshold:
                    keep[k] = False
        self.support_ = keep
        return self

    def transform(self, X):
        check_is_fitted(self, "support_")
        return self._subset(X)


class MutualInfoSelector(_FrameAwareSelector):
    """Keep the top-k features by mutual information with the target.

    Mutual information rather than a linear correlation because several
    relationships here are monotone but not linear -- cycle life against
    resistance is a clear case. Requires ``y``, so it can only ever be fitted
    where labels are legitimately available, i.e. the training fold.
    """

    def __init__(self, k: int | None = None, n_neighbors: int = 3, random_state: int = 42):
        self.k = k
        self.n_neighbors = n_neighbors
        self.random_state = random_state

    def fit(self, X, y=None):
        names = self._record_names(X)
        if y is None:
            raise ValueError("MutualInfoSelector requires y; it cannot be fitted unsupervised.")
        frame = pd.DataFrame(X, columns=names)
        # MI cannot consume NaN. Impute with training-fold medians for the
        # ranking only; the returned mask is applied to the untouched data.
        filled = frame.fillna(frame.median(numeric_only=True)).fillna(0.0)

        # The selectors serve both project routes, so the estimator must match
        # the target type: mutual_info_regression cannot consume the string
        # grade labels the secondary classification route supplies.
        target, is_categorical = _encode_target(y)
        estimator = mutual_info_classif if is_categorical else mutual_info_regression
        scores = estimator(
            filled, target.astype("int64") if is_categorical else target,
            n_neighbors=self.n_neighbors, random_state=self.random_state,
        )
        self.scores_ = scores
        k = len(names) if self.k is None else min(self.k, len(names))
        order = np.argsort(scores)[::-1][:k]
        self.support_ = np.zeros(len(names), dtype=bool)
        self.support_[order] = True
        return self

    def transform(self, X):
        check_is_fitted(self, "support_")
        return self._subset(X)


class StabilitySelector(_FrameAwareSelector):
    """Keep features selected in a sufficient fraction of bootstrap resamples.

    WHY, at n = 124: a single fit will happily rank a noise feature highly by
    chance. Requiring a feature to survive repeated resampling of the TRAINING
    fold is a far stronger criterion and is the standard remedy for selection
    instability in the small-n regime. All resampling happens strictly within
    whatever data ``fit`` receives.
    """

    def __init__(
        self,
        n_bootstrap: int = 100,
        sample_fraction: float = 0.75,
        selection_frequency_threshold: float = 0.6,
        k_per_bootstrap: int | None = None,
        random_state: int = 42,
    ):
        self.n_bootstrap = n_bootstrap
        self.sample_fraction = sample_fraction
        self.selection_frequency_threshold = selection_frequency_threshold
        self.k_per_bootstrap = k_per_bootstrap
        self.random_state = random_state

    def fit(self, X, y=None):
        names = self._record_names(X)
        if y is None:
            raise ValueError("StabilitySelector requires y.")
        frame = pd.DataFrame(X, columns=names)
        filled = frame.fillna(frame.median(numeric_only=True)).fillna(0.0)
        # Grade labels are encoded to their ordinal codes so the rank statistic
        # below is meaningful; see _encode_target.
        target, _ = _encode_target(y)

        rng = np.random.default_rng(self.random_state)
        n_samples = len(filled)
        draw = max(5, int(round(self.sample_fraction * n_samples)))
        k = self.k_per_bootstrap or max(1, len(names) // 2)

        counts = np.zeros(len(names), dtype=float)
        for _ in range(self.n_bootstrap):
            idx = rng.choice(n_samples, size=draw, replace=False)
            sub, sub_y = filled.iloc[idx], target[idx]
            if len(np.unique(sub_y)) < 3:
                continue
            # Spearman |rho| as the per-bootstrap ranking statistic: robust to
            # the monotone-but-nonlinear relationships this corpus contains.
            #
            # Computed VECTORISED over all features at once, and on NUMPY
            # arrays. Both details matter:
            #
            # * Vectorised because the loop version issued one scipy call per
            #   feature per bootstrap -- 5,700 calls per fit -- and dominated
            #   the Phase 6 sweep runtime. Spearman is Pearson on ranks, so
            #   ranking each column once per bootstrap and taking a single
            #   matrix product is exactly equivalent, not an approximation.
            # * On arrays because pandas correlates on the INDEX: the bootstrap
            #   subset keeps the original row labels while the target gets a
            #   fresh RangeIndex, so ``Series.corr`` silently mis-pairs them.
            #   That bug scored dq_var at |rho| = 0.02 against a true 0.91,
            #   demoting the strongest predictor in the project to noise.
            scores = _abs_spearman_matrix(sub.to_numpy(dtype="float64"), sub_y)
            counts[np.argsort(scores)[::-1][:k]] += 1.0

        self.selection_frequency_ = counts / max(self.n_bootstrap, 1)
        self.support_ = self.selection_frequency_ >= self.selection_frequency_threshold
        if not self.support_.any():
            # Never return zero features: fall back to the most stable ones.
            top = np.argsort(self.selection_frequency_)[::-1][:max(1, len(names) // 4)]
            self.support_ = np.zeros(len(names), dtype=bool)
            self.support_[top] = True
        return self

    def transform(self, X):
        check_is_fitted(self, "support_")
        return self._subset(X)


def build_selection_pipeline(bundle: ConfigBundle, *, include_scaler: bool = True) -> Pipeline:
    """Compose the configured preprocessing and selection steps into a Pipeline.

    The returned object is intended to be the front of a model pipeline, so that
    imputation, scaling and selection are all fitted on the training fold only.

    Ordering rationale: variance and correlation filters are unsupervised and
    run first to shrink the problem cheaply; the supervised steps (mutual
    information, stability selection) run afterwards on the reduced set.
    Imputation precedes them because the supervised selectors cannot consume
    NaN, and it uses the median, which is robust to the residual outliers the
    plausibility mask does not catch.
    """
    selection = bundle.features.selection
    steps: list[tuple[str, object]] = [
        # keep_empty_features=True is load-bearing, not a default tweak. Without
        # it SimpleImputer silently DROPS all-NaN columns, so the matrix leaving
        # this step has fewer columns than the feature registry declares and
        # every downstream selector's feature names shift out of alignment.
        # Keeping the column (filled with zeros) lets the variance filter remove
        # it explicitly, per fold, which is where that decision belongs.
        ("impute", SimpleImputer(strategy=bundle.features.preprocessing.impute_strategy,
                                 keep_empty_features=True)),
        ("variance", VarianceFilter(threshold=selection.variance_threshold)),
        ("correlation", CorrelationFilter(threshold=selection.correlation_threshold)),
    ]

    if selection.stability_selection.enabled:
        stability = selection.stability_selection
        steps.append((
            "stability",
            StabilitySelector(
                n_bootstrap=stability.n_bootstrap,
                sample_fraction=stability.sample_fraction,
                selection_frequency_threshold=stability.selection_frequency_threshold,
                random_state=bundle.models.seeds.global_,
            ),
        ))

    if selection.mutual_information.enabled:
        steps.append((
            "mutual_info",
            MutualInfoSelector(
                n_neighbors=selection.mutual_information.n_neighbors,
                random_state=bundle.models.seeds.global_,
            ),
        ))

    if include_scaler and bundle.features.preprocessing.scaler == "standard":
        steps.append(("scale", StandardScaler()))

    pipeline = Pipeline(steps)
    # Emit DataFrames between steps so feature names survive the whole chain.
    # Without this, selectors after the first receive a bare ndarray and fall
    # back to positional names, which would make the Phase 9 SHAP attributions
    # impossible to map back to a physical quantity.
    return pipeline.set_output(transform="pandas")