"""
Pandas implementation for dimensionality reduction feature groups.
"""

from __future__ import annotations

from typing import Any, TYPE_CHECKING, cast

from mloda.provider import ComputeFramework, DefaultOptionKeys
from mloda_plugins.compute_framework.base_implementations.pandas.dataframe import PandasDataFrame
from mloda_plugins.feature_group.experimental.dimensionality_reduction.base import DimensionalityReductionFeatureGroup

if TYPE_CHECKING:
    from numpy.typing import NDArray

try:
    import pandas as pd
    import numpy as np
except ImportError:
    pd = None
    np = None  # type: ignore[assignment]

# Check if required packages are available
SKLEARN_AVAILABLE = True
try:
    from sklearn.decomposition import PCA, FastICA
    from sklearn.manifold import TSNE, Isomap
    from sklearn.discriminant_analysis import LinearDiscriminantAnalysis
    from sklearn.preprocessing import StandardScaler
except ImportError:
    SKLEARN_AVAILABLE = False


class PandasDimensionalityReductionFeatureGroup(DimensionalityReductionFeatureGroup):
    @classmethod
    def compute_framework_rule(cls) -> set[type[ComputeFramework]]:
        """Define the compute framework for this feature group."""
        return {PandasDataFrame}

    @classmethod
    def _check_source_features_exist(cls, data: pd.DataFrame, feature_names: list[str]) -> None:
        """
        Check if the source features exist in the DataFrame.

        Args:
            data: The pandas DataFrame
            feature_names: List of feature names to check

        Raises:
            ValueError: If a feature does not exist in the DataFrame
        """
        missing_features = [f for f in feature_names if f not in data.columns]
        if missing_features:
            raise ValueError(
                f"Source features not found in data: {missing_features}. Available columns: {list(data.columns)}"
            )

    @classmethod
    def _add_result_to_data(cls, data: "pd.DataFrame", feature_name: str, result: "NDArray[Any]") -> "pd.DataFrame":
        """
        Add the dimensionality reduction result to the DataFrame using the multiple result columns pattern.

        Instead of storing the entire result array in a single column, this method creates multiple
        columns following the naming convention `feature_name~dim{i+1}` for each dimension.

        Args:
            data: The pandas DataFrame
            feature_name: The name of the feature to add
            result: The dimensionality reduction result (reduced features)

        Returns:
            The updated DataFrame with the dimensionality reduction result added as multiple columns
        """
        # Add individual dimension columns using the multiple result columns pattern
        algorithm, dimension = cls.parse_reduction_suffix(feature_name)
        named_columns = cls.apply_naming_convention(result, feature_name, suffix_generator=lambda i: f"dim{i + 1}")
        for col_name, col_data in named_columns.items():
            data[col_name] = col_data

        return data

    @classmethod
    def _perform_reduction(
        cls,
        data: Any,
        algorithm: str,
        dimension: int,
        source_features: list[str],
        options: Any,
    ) -> "NDArray[Any]":
        """
        Perform dimensionality reduction on the specified features.

        Args:
            data: The pandas DataFrame
            algorithm: The dimensionality reduction algorithm to use
            dimension: The target dimension for the reduction
            source_features: The list of source features to use for dimensionality reduction
            options: Options containing algorithm-specific parameters

        Returns:
            A numpy array containing the reduced features
        """
        # Import the base class to access option keys
        from mloda_plugins.feature_group.experimental.dimensionality_reduction.base import (
            DimensionalityReductionFeatureGroup,
        )

        # Cast data to pandas DataFrame
        df = cast(pd.DataFrame, data)

        # Extract the features to use for dimensionality reduction
        X = df[source_features].copy()

        # Handle missing values (replace with mean)
        for col in X.columns:
            if X[col].isna().any():
                X[col] = X[col].fillna(X[col].mean())

        # Convert to numpy array
        X_array = X.values

        # Check if scikit-learn is available
        if not SKLEARN_AVAILABLE:
            raise ImportError(
                "scikit-learn is required for dimensionality reduction. "
                "Please install it with 'pip install scikit-learn'."
            )

        # Standardize the features
        scaler = StandardScaler()
        X_scaled = scaler.fit_transform(X_array)

        # Check if dimension is valid
        if dimension >= X_scaled.shape[1]:
            raise ValueError(
                f"Target dimension ({dimension}) must be less than the number of source features ({X_scaled.shape[1]})"
            )

        # Perform dimensionality reduction based on the algorithm
        if algorithm == "pca":
            svd_solver = options.get(DimensionalityReductionFeatureGroup.PCA_SVD_SOLVER)
            if svd_solver is None:
                svd_solver = DimensionalityReductionFeatureGroup.PROPERTY_MAPPING[
                    DimensionalityReductionFeatureGroup.PCA_SVD_SOLVER
                ][DefaultOptionKeys.default]
            return cls._perform_pca_reduction(X_scaled, dimension, svd_solver)
        elif algorithm == "tsne":
            max_iter_val = options.get(DimensionalityReductionFeatureGroup.TSNE_MAX_ITER)
            if max_iter_val is None:
                max_iter_val = DimensionalityReductionFeatureGroup.PROPERTY_MAPPING[
                    DimensionalityReductionFeatureGroup.TSNE_MAX_ITER
                ][DefaultOptionKeys.default]
            max_iter = int(max_iter_val)

            n_iter_without_progress_val = options.get(DimensionalityReductionFeatureGroup.TSNE_N_ITER_WITHOUT_PROGRESS)
            if n_iter_without_progress_val is None:
                n_iter_without_progress_val = DimensionalityReductionFeatureGroup.PROPERTY_MAPPING[
                    DimensionalityReductionFeatureGroup.TSNE_N_ITER_WITHOUT_PROGRESS
                ][DefaultOptionKeys.default]
            n_iter_without_progress = int(n_iter_without_progress_val)

            method = options.get(DimensionalityReductionFeatureGroup.TSNE_METHOD)
            if method is None:
                method = DimensionalityReductionFeatureGroup.PROPERTY_MAPPING[
                    DimensionalityReductionFeatureGroup.TSNE_METHOD
                ][DefaultOptionKeys.default]
            return cls._perform_tsne_reduction(X_scaled, dimension, max_iter, n_iter_without_progress, method)
        elif algorithm == "ica":
            max_iter_val = options.get(DimensionalityReductionFeatureGroup.ICA_MAX_ITER)
            if max_iter_val is None:
                max_iter_val = DimensionalityReductionFeatureGroup.PROPERTY_MAPPING[
                    DimensionalityReductionFeatureGroup.ICA_MAX_ITER
                ][DefaultOptionKeys.default]
            max_iter = int(max_iter_val)
            return cls._perform_ica_reduction(X_scaled, dimension, max_iter)
        elif algorithm == "lda":
            return cls._perform_lda_reduction(X_scaled, dimension, df)
        elif algorithm == "isomap":
            n_neighbors_val = options.get(DimensionalityReductionFeatureGroup.ISOMAP_N_NEIGHBORS)
            if n_neighbors_val is None:
                n_neighbors_val = DimensionalityReductionFeatureGroup.PROPERTY_MAPPING[
                    DimensionalityReductionFeatureGroup.ISOMAP_N_NEIGHBORS
                ][DefaultOptionKeys.default]
            n_neighbors = int(n_neighbors_val)
            return cls._perform_isomap_reduction(X_scaled, dimension, n_neighbors)
        else:
            raise ValueError(f"Unsupported dimensionality reduction algorithm: {algorithm}")

    @classmethod
    def _perform_pca_reduction(cls, X: "NDArray[Any]", dimension: int, svd_solver: str = "auto") -> "NDArray[Any]":
        """
        Perform Principal Component Analysis (PCA).

        Args:
            X: The feature matrix
            dimension: The target dimension
            svd_solver: SVD solver algorithm ('auto', 'full', 'arpack', 'randomized')

        Returns:
            A numpy array containing the reduced features
        """
        # Check if scikit-learn is available
        if not SKLEARN_AVAILABLE:
            raise ImportError("scikit-learn is required for PCA dimensionality reduction")

        # Perform PCA
        pca = PCA(n_components=dimension, random_state=42, svd_solver=svd_solver)
        return cast("NDArray[Any]", pca.fit_transform(X))

    @classmethod
    def _perform_tsne_reduction(
        cls,
        X: "NDArray[Any]",
        dimension: int,
        max_iter: int = 250,
        n_iter_without_progress: int = 50,
        method: str = "barnes_hut",
    ) -> "NDArray[Any]":
        """
        Perform t-Distributed Stochastic Neighbor Embedding (t-SNE).

        Args:
            X: The feature matrix
            dimension: The target dimension
            max_iter: Maximum number of iterations for optimization
            n_iter_without_progress: Maximum iterations without progress before early stopping
            method: Computation method ('barnes_hut' or 'exact')

        Returns:
            A numpy array containing the reduced features
        """
        # Check if scikit-learn is available
        if not SKLEARN_AVAILABLE:
            raise ImportError("scikit-learn is required for t-SNE dimensionality reduction")

        # Adjust perplexity based on sample size
        # t-SNE requires perplexity < n_samples
        n_samples = X.shape[0]
        perplexity = min(30, n_samples // 3)  # Default is 30, but we need to ensure it's less than n_samples

        # For very small datasets (< 20 samples), use exact method as it's faster
        # and more reliable for unit tests
        if n_samples < 20:
            actual_method = "exact"
            # Use minimal iterations for small datasets in unit tests
            actual_max_iter = min(max_iter, 250)
        else:
            actual_method = method
            actual_max_iter = max_iter

        # Perform t-SNE with configurable parameters
        tsne = TSNE(
            n_components=dimension,
            random_state=42,
            perplexity=perplexity,
            max_iter=actual_max_iter,
            n_iter_without_progress=n_iter_without_progress,
            method=actual_method,
        )
        return cast("NDArray[Any]", tsne.fit_transform(X))

    @classmethod
    def _perform_ica_reduction(cls, X: "NDArray[Any]", dimension: int, max_iter: int = 200) -> "NDArray[Any]":
        """
        Perform Independent Component Analysis (ICA).

        Args:
            X: The feature matrix
            dimension: The target dimension
            max_iter: Maximum number of iterations

        Returns:
            A numpy array containing the reduced features
        """
        # Check if scikit-learn is available
        if not SKLEARN_AVAILABLE:
            raise ImportError("scikit-learn is required for ICA dimensionality reduction")

        # For small datasets, increase tolerance and iterations for better convergence
        n_samples = X.shape[0]
        if n_samples < 50:
            # Small datasets may need more iterations and higher tolerance
            actual_max_iter = max(max_iter, 1000)
            tol = 0.01
        else:
            actual_max_iter = max_iter
            tol = 1e-4  # sklearn default

        # Perform ICA
        ica = FastICA(n_components=dimension, random_state=42, max_iter=actual_max_iter, tol=tol)
        return cast("NDArray[Any]", ica.fit_transform(X))

    @classmethod
    def _perform_lda_reduction(cls, X: "NDArray[Any]", dimension: int, df: "pd.DataFrame") -> "NDArray[Any]":
        """
        Perform Linear Discriminant Analysis (LDA).

        Args:
            X: The feature matrix
            dimension: The target dimension
            df: The original DataFrame (needed for target variable)

        Returns:
            A numpy array containing the reduced features
        """
        # Check if scikit-learn is available
        if not SKLEARN_AVAILABLE:
            raise ImportError("scikit-learn is required for LDA dimensionality reduction")

        # LDA requires a target variable
        # We'll try to find a categorical column in the DataFrame
        categorical_columns = df.select_dtypes(include=["object", "string", "category"]).columns

        if len(categorical_columns) == 0:
            raise ValueError("LDA requires a categorical target variable, but none was found in the data")

        # Use the first categorical column as the target
        target_column = categorical_columns[0]
        y = df[target_column].values

        # Perform LDA
        lda = LinearDiscriminantAnalysis(n_components=dimension)
        return cast("NDArray[Any]", lda.fit_transform(X, y))

    @classmethod
    def _perform_isomap_reduction(cls, X: "NDArray[Any]", dimension: int, n_neighbors: int = 5) -> "NDArray[Any]":
        """
        Perform Isometric Mapping (Isomap).

        Args:
            X: The feature matrix
            dimension: The target dimension
            n_neighbors: Number of neighbors to consider for each point

        Returns:
            A numpy array containing the reduced features
        """
        # Check if scikit-learn is available
        if not SKLEARN_AVAILABLE:
            raise ImportError("scikit-learn is required for Isomap dimensionality reduction")

        # Perform Isomap
        isomap = Isomap(n_components=dimension, n_neighbors=n_neighbors)
        return cast("NDArray[Any]", isomap.fit_transform(X))
