"""
Computes Mean Decrease Impurity (MDI) feature importance.
"""

import pandas as pd
import numpy as np
from typing import List, Optional, Union, Any
from sklearn.ensemble import BaseEnsemble
from .feature_importance_strategy import FeatureImportanceStrategy

class FeatureImportanceMDI(FeatureImportanceStrategy):
    """
    Computes feature importance using Mean Decrease Impurity (MDI).

    This method is specific to tree-based ensembles (like RandomForest)
    and measures importance as the average impurity decrease.
    """

    def __init__(self, classifier: BaseEnsemble):
        """
        Initialize the strategy.

        Parameters
        ----------
        classifier : BaseEnsemble
            An *untrained* scikit-learn ensemble model (e.g., RandomForestClassifier).
        """
        if not hasattr(classifier, 'estimators_'):
            raise TypeError("Classifier must be an ensemble (e.g., RandomForest).")
        self.classifier = classifier

    def compute(self, x: pd.DataFrame, y: pd.Series, **kwargs: Any) -> pd.DataFrame:
        """
        Compute MDI feature importance.

        Parameters
        ----------
        x : pd.DataFrame
            The feature data.
        y : pd.Series
            The target data.
        **kwargs : Any
            Keyword arguments for the classifier's `fit` method
            (e.g., `sample_weight`).

        Returns
        -------
        pd.DataFrame
            DataFrame with "Mean" and "StandardDeviation" of importance.
        """
        train_sample_weights = kwargs.get('sample_weight')
        
        # Fit the classifier
        self.classifier.fit(x, y, sample_weight=train_sample_weights)

        # Get importance from each tree
        importances_dict = {
            i: tree.feature_importances_
            for i, tree in enumerate(self.classifier.estimators_)
        }
        importances_df = pd.DataFrame.from_dict(importances_dict, orient="index")
        
        # Ensure correct feature names
        if hasattr(self.classifier, 'feature_names_in_'):
             importances_df.columns = self.classifier.feature_names_in_
        else:
             importances_df.columns = x.columns

        # Replace 0 with NaN (as per user's original code)
        importances_df.replace(0, np.nan, inplace=True)

        importances = pd.concat(
            {
                "Mean": importances_df.mean(),
                "StandardDeviation": (
                    importances_df.std()
                    * (importances_df.shape[0] ** -0.5)
                ),
            },
            axis=1,
        )

        # Normalize
        importances /= importances["Mean"].sum()
        return importances