from ..utils import alpha_feature_importance


class GlobalFeatureImportance:
    def get_alpha_feature_importance(self, alpha):

        feat_imp = alpha_feat_imp = self.feature_importance.set_index(
            "Variable"
        ).sort_values("Importance", ascending=False)

        if alpha is not None:
            alpha_feat_imp = alpha_feature_importance(feat_imp, alpha)

        alpha_cond_feat_imp = None
        if hasattr(self, "conditional_feature_importance") and (
            self.conditional_feature_importance is not None
        ):
            alpha_cond_feat_imp = {
                label: value.set_index("Variable")
                .sort_values("Importance", ascending=False)
                .iloc[: len(alpha_feat_imp.index)]
                for label, value in self.conditional_feature_importance.items()
            }

        return feat_imp, (alpha_feat_imp, alpha_cond_feat_imp)

    def partial_dependence_plot(self, first=0, last=None, **plot_kargs):

        import matplotlib.pyplot as plt
        from sklearn.inspection import PartialDependenceDisplay

        top_k = None
        if last == None:
            last = first + 6

        importances = self.get_alpha_feature_importance(alpha=top_k)
        feat_imp, (alpha_feat_imp, alpha_cond_feat_imp) = importances
        features = list(alpha_feat_imp.index)[first:last]
        title = "Partial dependence plot"
        percentiles = (
            (0, 1) if self.model_type == "binary_classification" else (0.05, 0.95)
        )

        common_params = {
            "subsample": 50,
            "n_jobs": 2,
            "grid_resolution": 20,
            "random_state": 0,
            "kind": "average",
            "percentiles": percentiles,
            "line_kw": {"color": "mediumslateblue", "label": "Average"},
        }

        common_params.update(plot_kargs)

        plt.rcParams["figure.constrained_layout.use"] = True

        pdp = PartialDependenceDisplay.from_estimator(
            self.model,
            self.x,
            features,
            **common_params,
        )
        pdp.figure_.suptitle(title)


class LocalFeatureImportance:
    pass


class BaseFeatureImportance:
    def __init__(
        self,
        model_type,
        model,
        x,
        y,
        importance_weights,
        conditional_importance_weights,
    ):
        self.model_type = model_type
        self.model = model
        self.x = x
        self.y = y
        self.importance_weights = importance_weights
        self.conditional_importance_weights = conditional_importance_weights

    def custom_metrics(self):
        pass
