"""
Feature selection pipeline: correlation analysis и Top-K feature selection.
"""

import pandas as pd
import numpy as np
from sklearn.feature_selection import SelectKBest, mutual_info_classif
from utils.logger import logger


class FeatureSelector:
    """
    Пайплайн для отбора лучших фичей с использованием Mutual Information.

    Args:
        k: Количество фичей для сохранения

    Attributes:
        k: Количество фичей
        selector: sklearn SelectKBest объект
        selected_features: Список отобранных фичей

    Example:
        >>> selector = FeatureSelector(k=30)
        >>> X_clean = selector.fit_transform(X, y)
        >>> print(selector.selected_features)
    """

    def __init__(self, k: int = 30):
        """
        Инициализация FeatureSelector.

        Args:
            k: Количество фичей для отбора (по умолчанию 30)
        """
        self.k = k
        self.selector = None
        self.selected_features = None
        self.scores = None

    def fit(self, X: pd.DataFrame, y: np.ndarray):
        """
        Обучает селектор на данных.

        Args:
            X: Feature matrix (n_samples, n_features)
            y: Target vector (n_samples,) or (n_samples, n_targets)
        """
        logger.info(
            f"Running SelectKBest (mutual_info) to keep top {min(self.k, X.shape[1])} features..."
        )

        # Mutual Information требует 1D массив для y
        # Используем первый таргет (Long) для оценки важности фичей
        if len(y.shape) > 1:
            logger.info("Using 'Long' target for feature importance evaluation")
            y_1d = y[:, 0]
        else:
            y_1d = y

        self.selector = SelectKBest(
            score_func=mutual_info_classif, k=min(self.k, X.shape[1])
        )
        self.selector.fit(X, y_1d)

        # Получаем скоры и выбираем топ-K
        self.scores = pd.Series(self.selector.scores_, index=X.columns).sort_values(
            ascending=False
        )
        self.selected_features = self.scores.head(self.k).index.tolist()

        logger.info(f"Top {len(self.selected_features)} features selected.")
        if self.k < len(X.columns):
            logger.info(f"Removed {len(X.columns) - self.k} low-importance features.")

        return self

    def transform(self, X: pd.DataFrame) -> pd.DataFrame:
        """
        Применяет отбор фичей к новым данным.

        Args:
            X: Feature matrix (n_samples, n_features)

        Returns:
            pd.DataFrame: Feature matrix с отобранными фичами (n_samples, k)
        """
        if self.selected_features is None:
            raise RuntimeError(
                "Selector is not fitted yet. Call fit() first."
            )
        return X[self.selected_features]

    def fit_transform(self, X: pd.DataFrame, y: np.ndarray) -> pd.DataFrame:
        """
        Обучает и применяет отбор фичей.

        Args:
            X: Feature matrix (n_samples, n_features)
            y: Target vector (n_samples,) or (n_samples, n_targets)

        Returns:
            pd.DataFrame: Feature matrix с отобранными фичами
        """
        self.fit(X, y)
        return self.transform(X)

    def get_feature_scores(self) -> pd.Series:
        """
        Возвращает скоры важности для всех фичей.

        Returns:
            pd.Series: Скоры важности по убыванию
        """
        if self.scores is None:
            raise RuntimeError("Fit first to get feature scores.")
        return self.scores

    def get_selected_feature_names(self) -> list[str]:
        """
        Возвращает имена отобранных фичей.

        Returns:
            list[str]: Имена топ-K фичей
        """
        return self.selected_features

    def is_fitted(self) -> bool:
        """Возвращает True, если селектор обучен."""
        return self.selected_features is not None