"""
Base implementation for missing value imputation feature groups.
"""

from __future__ import annotations

import copy
from abc import abstractmethod
from typing import Any, List, Optional, Set, Union

from mloda import FeatureGroup
from mloda import Feature
from mloda.provider import FeatureChainParser
from mloda.provider import (
    FeatureChainParserMixin,
)
from mloda.user import FeatureName
from mloda.provider import FeatureSet
from mloda import Options
from mloda_plugins.feature_group.experimental.default_options_key import DefaultOptionKeys


class MissingValueFeatureGroup(FeatureChainParserMixin, FeatureGroup):
    """
    Base class for all missing value imputation feature groups.

    Missing value feature groups impute missing values in the source feature using
    the specified imputation method. They support both string-based feature creation
    and configuration-based creation with proper group/context parameter separation.

    ## Supported Imputation Methods

    - `mean`: Impute with the mean of non-missing values
    - `median`: Impute with the median of non-missing values
    - `mode`: Impute with the most frequent value
    - `constant`: Impute with a specified constant value
    - `ffill`: Forward fill (use the last valid value)
    - `bfill`: Backward fill (use the next valid value)

    ## Feature Creation Methods

    ### 1. String-Based Creation

    Features follow the naming pattern: `{in_features}__{imputation_method}_imputed`

    Examples:
    ```python
    features = [
        "income__mean_imputed",      # Impute missing values in income with the mean
        "age__median_imputed",       # Impute missing values in age with the median
        "category__constant_imputed" # Impute missing values in category with a constant value
    ]
    ```

    ### 2. Configuration-Based Creation

    Uses Options with proper group/context parameter separation:

    ```python
    feature = Feature(
        name="placeholder",  # Placeholder name, will be replaced
        options=Options(
            context={
                MissingValueFeatureGroup.IMPUTATION_METHOD: "mean",
                DefaultOptionKeys.in_features: "income",
            }
        )
    )
    ```

    ## Parameter Classification

    ### Context Parameters (Default)
    These parameters don't affect Feature Group resolution/splitting:
    - `imputation_method`: The type of imputation to perform
    - `in_features`: The source feature to impute missing values
    - `constant_value`: Constant value for constant imputation (optional)
    - `group_by_features`: Features to group by before imputation (optional)

    ### Group Parameters
    Currently none for MissingValueFeatureGroup. Parameters that affect Feature Group
    resolution/splitting would be placed here.

    ## Usage Examples

    ### String-Based Creation

    ```python
    from mloda import Feature

    # Impute missing income values with mean
    feature = Feature(name="income__mean_imputed")

    # Impute missing age values with median
    feature = Feature(name="age__median_imputed")

    # Impute missing category values with mode
    feature = Feature(name="category__mode_imputed")

    # Forward fill missing temperature values
    feature = Feature(name="temperature__ffill_imputed")
    ```

    ### Configuration-Based Creation

    ```python
    from mloda import Feature
    from mloda import Options
    from mloda_plugins.feature_group.experimental.default_options_key import DefaultOptionKeys

    # Mean imputation using configuration
    feature = Feature(
        name="placeholder",
        options=Options(
            context={
                MissingValueFeatureGroup.IMPUTATION_METHOD: "mean",
                DefaultOptionKeys.in_features: "income",
            }
        )
    )

    # Constant imputation with a specific value
    feature = Feature(
        name="placeholder",
        options=Options(
            context={
                MissingValueFeatureGroup.IMPUTATION_METHOD: "constant",
                DefaultOptionKeys.in_features: "status",
                "constant_value": "unknown",
            }
        )
    )

    # Group-based imputation (e.g., mean by category)
    feature = Feature(
        name="placeholder",
        options=Options(
            context={
                MissingValueFeatureGroup.IMPUTATION_METHOD: "mean",
                DefaultOptionKeys.in_features: "price",
                "group_by_features": ["product_category", "region"],
            }
        )
    )
    ```

    ## Requirements
    - Input data must contain the source feature to be imputed
    - For group-based imputation, grouping features must also be present
    - For constant imputation, a constant_value must be provided
    """

    IMPUTATION_METHOD = "imputation_method"
    # Define supported imputation methods
    IMPUTATION_METHODS = {
        "mean": "Impute with the mean of non-missing values",
        "median": "Impute with the median of non-missing values",
        "mode": "Impute with the most frequent value",
        "constant": "Impute with a specified constant value",
        "ffill": "Forward fill (use the last valid value)",
        "bfill": "Backward fill (use the next valid value)",
    }

    PREFIX_PATTERN = r".*__([\w]+)_imputed$"

    # In-feature configuration for FeatureChainParserMixin
    MIN_IN_FEATURES = 1
    MAX_IN_FEATURES = 1

    PROPERTY_MAPPING = {
        IMPUTATION_METHOD: {
            **IMPUTATION_METHODS,
            DefaultOptionKeys.context: True,
        },
        DefaultOptionKeys.in_features: {
            "explanation": "Source feature to impute missing values",
            DefaultOptionKeys.context: True,
        },
        "constant_value": {
            "explanation": "Constant value to use for constant imputation method",
            DefaultOptionKeys.context: True,
            DefaultOptionKeys.default: None,  # Default is None, required only for constant method
        },
        "group_by_features": {
            "explanation": "Optional list of features to group by before imputation",
            DefaultOptionKeys.context: True,
            DefaultOptionKeys.default: None,  # Default is None (no grouping)
        },
    }

    @classmethod
    def get_imputation_method(cls, feature_name: str) -> str:
        """Extract the imputation method from the feature name."""
        # parse_feature_name returns (operation_config, source_feature)
        # The operation_config contains the imputation method extracted from the suffix pattern
        operation_config, _ = FeatureChainParser.parse_feature_name(feature_name, [cls.PREFIX_PATTERN])
        if operation_config is None:
            raise ValueError(f"Invalid missing value feature name format: {feature_name}")

        # The PREFIX_PATTERN captures the method name (e.g., "mean" from "mean_imputed")
        # So operation_config already contains just the method name
        imputation_method = operation_config

        # Validate imputation method
        if imputation_method not in cls.IMPUTATION_METHODS:
            raise ValueError(
                f"Unsupported imputation method: {imputation_method}. "
                f"Supported methods: {', '.join(cls.IMPUTATION_METHODS.keys())}"
            )

        return imputation_method

    @classmethod
    def _extract_imputation_method(cls, feature: Feature) -> Optional[str]:
        """
        Extract imputation method from a feature.

        Tries string-based parsing first, falls back to configuration-based.

        Args:
            feature: The feature to extract imputation method from

        Returns:
            Imputation method name or None if not found
        """
        feature_name = feature.get_name()

        # Try string-based parsing first
        if FeatureChainParser.is_chained_feature(feature_name):
            # Use get_imputation_method which handles parse_feature_name correctly
            return cls.get_imputation_method(feature_name)

        # Fall back to configuration-based approach
        imputation_method = feature.options.get(cls.IMPUTATION_METHOD)

        # Validate imputation method if found
        if imputation_method is not None and imputation_method not in cls.IMPUTATION_METHODS:
            raise ValueError(
                f"Unsupported imputation method: {imputation_method}. "
                f"Supported methods: {', '.join(cls.IMPUTATION_METHODS.keys())}"
            )

        return str(imputation_method) if imputation_method is not None else None

    @classmethod
    def _extract_imputation_method_and_source_feature(cls, feature: Feature) -> tuple[str, str]:
        """
        Extract imputation method and source feature name from a feature.

        Tries string-based parsing first, falls back to configuration-based approach.

        Args:
            feature: The feature to extract parameters from

        Returns:
            Tuple of (imputation_method, source_feature_name)

        Raises:
            ValueError: If parameters cannot be extracted
        """
        source_features = cls._extract_source_features(feature)
        imputation_method = cls._extract_imputation_method(feature)

        if imputation_method is None:
            raise ValueError(f"Could not extract imputation method from: {feature.name}")

        return imputation_method, source_features[0]

    @classmethod
    def calculate_feature(cls, data: Any, features: FeatureSet) -> Any:
        """
        Perform missing value imputation.

        Processes all requested features, determining the imputation method
        and source feature from either string parsing or configuration-based options.

        Adds the imputed results directly to the input data structure.
        """

        # Process each requested feature
        for feature in features.features:
            imputation_method, source_feature = cls._extract_imputation_method_and_source_feature(feature)

            # Resolve multi-column features automatically
            # If source_feature is "onehot_encoded__product", this discovers
            # ["onehot_encoded__product~0", "onehot_encoded__product~1", ...]
            available_columns = cls._get_available_columns(data)
            resolved_columns = cls.resolve_multi_column_feature(source_feature, available_columns)

            constant_value = feature.options.get("constant_value")
            group_by_features = feature.options.get("group_by_features")

            cls._check_source_features_exist(data, resolved_columns)

            # Validate group by features if provided
            if group_by_features:
                for group_feature in group_by_features:
                    cls._check_source_features_exist(data, [group_feature])

            # Validate constant value is provided for constant imputation
            if imputation_method == "constant" and constant_value is None:
                raise ValueError("Constant value must be provided for constant imputation method")

            # Apply the appropriate imputation function
            result = cls._perform_imputation(
                data, imputation_method, resolved_columns, constant_value, group_by_features
            )

            # Add the result to the data
            data = cls._add_result_to_data(data, feature.get_name(), result)
        return data

    @classmethod
    @abstractmethod
    def _get_available_columns(cls, data: Any) -> Set[str]:
        """
        Get the set of available column names from the data.

        Args:
            data: The input data

        Returns:
            Set of column names available in the data
        """
        ...

    @classmethod
    @abstractmethod
    def _check_source_features_exist(cls, data: Any, feature_names: List[str]) -> None:
        """
        Check if the resolved source features exist in the data.

        Args:
            data: The input data
            feature_names: List of resolved feature names (may contain ~N suffixes)

        Raises:
            ValueError: If none of the features exist in the data
        """
        ...

    @classmethod
    @abstractmethod
    def _add_result_to_data(cls, data: Any, feature_name: str, result: Any) -> Any:
        """
        Add the result to the data.

        Args:
            data: The input data
            feature_name: The name of the feature to add
            result: The result to add

        Returns:
            The updated data
        """
        ...

    @classmethod
    @abstractmethod
    def _perform_imputation(
        cls,
        data: Any,
        imputation_method: str,
        in_features: List[str],
        constant_value: Optional[Any] = None,
        group_by_features: Optional[List[str]] = None,
    ) -> Any:
        """
        Method to perform the imputation. Should be implemented by subclasses.

        Supports both single-column and multi-column imputation:
        - Single column: [feature_name] - imputes values within the column
        - Multi-column: [feature~0, feature~1, ...] - imputes across columns

        Args:
            data: The input data
            imputation_method: The type of imputation to perform
            in_features: List of resolved source feature names to impute
            constant_value: The constant value to use for imputation (if method is 'constant')
            group_by_features: Optional list of features to group by before imputation

        Returns:
            The result of the imputation
        """
        ...
