# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/error_analysis.error_analysis.ipynb.

# %% auto 0
__all__ = ['ErrorAnalyzer']

# %% ../../nbs/error_analysis.error_analysis.ipynb 3
from typing import Optional, Sequence, Union, Dict

import numpy as np
import pandas as pd
from sklearn.inspection import permutation_importance
from sklearn.base import BaseEstimator

from ..preprocessing import PoniardPreprocessor
from ..utils.utils import get_kwargs, non_default_repr
from ..utils.estimate import element_to_list_maybe, get_target_info
from ..estimators.core import PoniardBaseEstimator

# %% ../../nbs/error_analysis.error_analysis.ipynb 5
class ErrorAnalyzer:
    """An error analyzer for predictive models.

    Compare ground truth and predicted target and rank the largest deviations
    (either by probabilities for classifiers and actual values for regressors).

    This class is tightly integrated with `PoniardBaseEstimator`, but does not require it.

    Parameters
    ----------
    task :
        The machine learning task. Either 'regression' or 'classification'.
    """

    def __init__(self, task: str):
        self._init_params = get_kwargs()
        self.task = task
        self._poniard: Optional["PoniardBaseEstimator"] = None

    @property
    def _has_poniard(self):
        return True if self._poniard is not None else False

    @classmethod
    def from_poniard(
        cls, poniard: "PoniardBaseEstimator", estimator_names: Union[str, Sequence[str]]
    ):
        """Use a Poniard instance to instantiate `ErrorAnalyzer`.

        Automatically sets the task and gives access to the underlying data.

        Parameters
        ----------
        poniard :
            A `PoniardClassifier` or `PoniardRegressor` instance.
        estimator_names :
            Array of estimators for which to compute errors.

        Returns
        -------
        ErrorAnalyzer :
            An instance of the class.
        """
        error_analysis = cls(task=poniard.poniard_task)
        error_analysis._poniard = poniard
        error_analysis.estimator_names = element_to_list_maybe(estimator_names)
        error_analysis.type_of_target = poniard.target_info["type_"]
        return error_analysis

    def _compute_predictions(self):
        """Compute predictions for Poniard estimators."""
        predictions = self._poniard.predict(estimator_names=self.estimator_names)
        probas = None
        if self.type_of_target in ["binary", "multilabel-indicator", "multiclass"]:
            probas = self._poniard.predict_proba(estimator_names=self.estimator_names)
        return predictions, probas

    def rank_errors(
        self,
        y: Optional[Union[np.ndarray, pd.Series, pd.DataFrame]] = None,
        predictions: Optional[Union[np.ndarray, pd.Series, pd.DataFrame]] = None,
        probas: Optional[Union[np.ndarray, pd.Series, pd.DataFrame]] = None,
        exclude_correct: bool = True,
    ):
        """Compare the `y` ground truth with `predictions` and `probas` and sort by the largest deviations.

        If `ErrorAnalyzer.from_poniard` was used, no data needs to be passed to this method.

        In this context "error" refers to:

        * misclassified samples in binary and multiclass problems.
        * misclassified samples in any of the labels for multilabel problems.
        * samples with predicted values outside the `truth - 1SD <-> truth + 1SD`
        range for regression problems.
        * samples with predicted values outside the `truth - 1SD <-> truth + 1SD`
        range in any of the targets for multioutput regression problems.

        Parameters
        ----------
        y :
            Ground truth target.
        predictions :
            Predicted target.
        probas :
            Predicted probabilities for each class in classification tasks.
        exclude_correct :
            Whether to exclude correctly predicted samples in the output ranking. Default True.

        Returns
        -------
        Dict
            Ranked errors
        """
        if self._has_poniard:
            y = self._poniard.y
            predictions, probas = self._compute_predictions()
            type_of_target = self.type_of_target
            ranked_errors = {}
            for estimator in self.estimator_names:
                proc_probas = probas[estimator] if probas is not None else probas
                estimator_errors = self._target_redirect(type_of_target)(
                    y, predictions[estimator], proc_probas, exclude_correct
                )
                ranked_errors.update({estimator: estimator_errors})
            return ranked_errors
        else:
            self.type_of_target = get_target_info(y, task=self.task)["type_"]
            return self._target_redirect(self.type_of_target)(
                y, predictions, probas, exclude_correct
            )

    def _target_redirect(self, type_of_target: str):
        """A router for error ranking depending on the type of the target."""
        if type_of_target == "binary":
            return self._rank_errors_binary
        elif type_of_target == "multiclass":
            return self._rank_errors_multiclass
        elif type_of_target == "multilabel-indicator":
            return self._rank_errors_multilabel
        elif type_of_target == "continuous":
            return self._rank_errors_continuous
        elif type_of_target == "continuous-multioutput":
            return self._rank_errors_continuous_multioutput
        else:
            raise NotImplementedError("Type of target could not be determined.")

    def _rank_errors_binary(
        self,
        y: Union[np.ndarray, pd.Series, pd.DataFrame],
        predictions: Union[np.ndarray, pd.Series, pd.DataFrame],
        probas: Union[np.ndarray, pd.Series, pd.DataFrame],
        exclude_correct: bool = True,
    ):
        errors = pd.DataFrame(
            {
                "y": y,
                "prediction": predictions,
                "proba_0": probas[:, 0],
                "proba_1": probas[:, 1],
            }
        )
        if exclude_correct:
            errors = errors.query("y != prediction")
        errors = errors.assign(error=(errors["y"] - errors["proba_1"]).abs())
        ranked_errors = errors.sort_values("error", ascending=False)
        return {"values": ranked_errors, "idx": ranked_errors.index}

    def _rank_errors_multiclass(
        self,
        y: Union[np.ndarray, pd.Series, pd.DataFrame],
        predictions: Union[np.ndarray, pd.Series, pd.DataFrame],
        probas: Union[np.ndarray, pd.Series, pd.DataFrame],
        exclude_correct: bool = True,
    ):
        data = {"y": y, "prediction": predictions}
        data.update({f"proba_{i}": probas[:, i] for i in range(len(np.unique(y)))})
        errors = pd.DataFrame(data)
        if exclude_correct:
            errors = errors.query("y != prediction")
        errors = errors.assign(
            truth_proba=[
                errors[f"proba_{truth}"].iloc[idx]
                for idx, truth in enumerate(errors["y"])
            ]
        )
        errors = errors.assign(error=(1 - errors["truth_proba"]).abs())
        ranked_errors = errors.sort_values("error", ascending=False)
        return {"values": ranked_errors, "idx": ranked_errors.index}

    def _rank_errors_multilabel(
        self,
        y: Union[np.ndarray, pd.Series, pd.DataFrame],
        predictions: Union[np.ndarray, pd.Series, pd.DataFrame],
        probas: Union[np.ndarray, pd.Series, pd.DataFrame],
        exclude_correct: bool = True,
    ):
        truth = pd.DataFrame(y, columns=[f"y_{i}" for i in range(y.shape[1])])
        preds = pd.DataFrame(
            predictions,
            columns=[f"prediction_{i}" for i in range(y.shape[1])],
        )
        pro = pd.DataFrame(probas, columns=[f"proba_{i}" for i in range(y.shape[1])])
        errors = pd.concat([truth, preds, pro], axis=1)
        if exclude_correct:
            errors = errors.loc[~preds.eq(y).all(axis=1)]
        errors_per_label = (
            errors[truth.columns].sub(errors[pro.columns].values, axis=1)
        ).abs()
        last = lambda x: x[-1]
        zero_array = np.zeros_like(errors_per_label)
        # It might be reasonable to avoid considering errors if the ground truth aligns with the proba.
        # errors_per_label = errors_per_label.mask(
        #    errors[truth.columns]
        #    .rename(columns=last)
        #    .eq(errors[preds.columns].rename(columns=last))
        #    .values,
        #    zero_array,
        # )
        errors_per_label.columns = [f"error_{i}" for i in range(y.shape[1])]
        errors_per_label = errors_per_label.assign(error=errors_per_label.mean(axis=1))
        errors = pd.concat([errors, errors_per_label], axis=1)
        ranked_errors = errors.sort_values("error", ascending=False)
        return {"values": ranked_errors, "idx": ranked_errors.index}

    def _rank_errors_continuous(
        self,
        y: Union[np.ndarray, pd.Series, pd.DataFrame],
        predictions: Union[np.ndarray, pd.Series, pd.DataFrame],
        probas=None,
        exclude_correct: bool = True,
    ):
        errors = pd.DataFrame({"y": y, "prediction": predictions})
        y_std = np.std(y)
        if exclude_correct:
            errors = errors.query("prediction < y - @y_std | prediction > y + @y_std")
        errors = errors.assign(error=(errors["y"] - errors["prediction"]).abs())
        ranked_errors = errors.sort_values("error", ascending=False)
        return {"values": ranked_errors, "idx": ranked_errors.index}

    def _rank_errors_continuous_multioutput(
        self,
        y: Union[np.ndarray, pd.Series, pd.DataFrame],
        predictions: Union[np.ndarray, pd.Series, pd.DataFrame],
        probas=None,
        exclude_correct: bool = True,
    ):
        truth = pd.DataFrame(y, columns=[f"y_{i}" for i in range(y.shape[1])])
        preds = pd.DataFrame(
            predictions,
            columns=[f"prediction_{i}" for i in range(y.shape[1])],
        )
        errors = pd.concat([truth, preds], axis=1)
        if exclude_correct:
            y_std = np.std(y, axis=0)
            errors = errors.loc[
                (
                    (preds.values < truth.values - y_std)
                    | (preds.values > truth.values + y_std)
                ).any(axis=1)
            ].loc[lambda x: ~x.index.duplicated()]
        errors_per_target = pd.DataFrame(
            np.abs(errors[truth.columns].values - errors[preds.columns].values)
        ).set_index(errors.index)
        errors_per_target.columns = [f"error_{i}" for i in range(y.shape[1])]
        errors_per_target = errors_per_target.assign(
            error=errors_per_target.mean(axis=1)
        )
        errors = pd.concat([errors, errors_per_target], axis=1)
        ranked_errors = errors.sort_values("error", ascending=False)
        return {"values": ranked_errors, "idx": ranked_errors.index}

    @staticmethod
    def merge_errors(errors: Dict[str, Dict[str, Union[pd.DataFrame, pd.Series]]]):
        """Merge multiple error rankings. This is particularly useful when evaluating multiple estimators.

        Compute how many estimators had the specific error and the average error between them.

        This method works best when using `ErrorAnalyzer.from_poniard`, since `errors` can be
        the output of `ErrorAnalyzer.rank_errors`. However, this is not required; as long as
        `errors` is properly defined (`{str: {str: pd.DataFrame, str: pd.Series}}`)

        Parameters
        ----------
        errors :
            Dictionary of errors and error indexes.

        Returns
        -------
        Dict
            Merged errors
        """
        concatenated = pd.concat(
            [
                error_dict["values"].assign(estimator=estimator)
                for estimator, error_dict in errors.items()
            ]
        ).reset_index()
        concatenated = (
            concatenated.groupby("index")
            .agg(
                mean_error=pd.NamedAgg(column="error", aggfunc=np.mean),
                freq=pd.NamedAgg(column="error", aggfunc=np.size),
                estimators=pd.NamedAgg(column="estimator", aggfunc=lambda x: list(x)),
            )
            .sort_values(["freq", "mean_error"], ascending=False)
        )
        return {"values": concatenated, "idx": concatenated.index}

    def analyze_target(
        self,
        errors_idx: pd.Series,
        y: Optional[Union[np.ndarray, pd.Series, pd.DataFrame]] = None,
        reg_bins: int = 5,
        as_ratio: bool = False,
        wrt_target: bool = False,
    ):
        """Analyze which target classes/ranges have the most errors and compare with observed
        target distribution.

        Parameters
        ----------
        errors_idx :
            Index of ranked errors.
        y :
            Ground truth. Not needed if using `ErrorAnalyzer.from_poniard`.
        reg_bins :
            Number of bins in which to place ground truth targets for regression tasks.
        as_ratio :
            Whether to show error ratios instead of error counts per class/bin. Default False.
        wrt_target :
            Whether to compute counts of errors or error ratios with respect
            to the ground truth. Default False.

        Returns
        -------
        pd.DataFrame
            Counts per error.
        """
        type_of_target = self.type_of_target
        if self._has_poniard:
            y = self._poniard.y
        y = pd.DataFrame(y)
        y_errors = y.loc[errors_idx]

        if type_of_target in ["binary", "multiclass", "multilabel-indicator"]:
            target_names = y.columns.tolist()
        elif type_of_target == "continuous":
            bins = pd.qcut(y.squeeze(), q=reg_bins)
            y = y.assign(bins=bins)
            y_errors = y_errors.assign(bins=bins)
            target_names = "bins"
        elif type_of_target == "continuous-multioutput":
            bins = {
                f"bin_{target}": pd.qcut(y[target], q=reg_bins)
                for target in range(y.shape[1])
            }
            y = y.assign(**bins)
            y_errors = y_errors.assign(**bins)
            target_names = list(bins.keys())
        else:
            raise NotImplementedError("Type of target could not be determined.")
        errors_dist = y_errors.groupby(target_names).size()
        target_dist = y.groupby(target_names).size()
        if as_ratio:
            errors_dist = errors_dist / errors_dist.sum()
            target_dist = target_dist / target_dist.sum()
        if wrt_target:
            output = (errors_dist / target_dist).fillna(0).sort_values(ascending=False)
        else:
            output = pd.DataFrame(errors_dist).join(
                pd.DataFrame(target_dist),
                how="right",
                lsuffix="_errors",
                rsuffix="_target",
            )
            output = output.fillna(0).sort_values(by=output.columns[0], ascending=False)
        return output

    def analyze_features(
        self,
        errors_idx: pd.Series,
        X: Optional[Union[np.ndarray, pd.Series, pd.DataFrame]] = None,
        features: Optional[Sequence[Union[str, int]]] = None,
        estimator_name: Optional[Union[str, BaseEstimator]] = None,
        n_features: Optional[Union[int, float]] = None,
    ):
        """Cross tabulate features with prediction errors.

        Parameters
        ----------
        errors_idx :
            Index of ranked errors.
        X :
            Features array. Not needed if using `ErrorAnalyzer.from_poniard`.
        features :
            Array of features to analyze. If `None`, all features will be analyzed.
        estimator_name :
            Only valid if using `ErrorAnalyzer.from_poniard`. Allows using an estimator to
            compute permutation importances and analyzing only the top `n_features`.
        n_features :
            How many features to analyze based on permutation importances.

        Returns
        -------
        Dict[str, pd.DataFrame]
            Per feature summary.
        """
        if self._has_poniard:
            X = self._poniard.X
            feature_types = self._poniard.feature_types.items()
        else:
            feature_types = (
                PoniardPreprocessor(task="placeholder")
                .build(X, np.zeros((X.shape[0],)))
                .feature_types
            )
        inverted_feature_types = {}
        for k, v in feature_types:
            for i in v:
                inverted_feature_types[i] = k
        X = pd.DataFrame(X).assign(error=lambda x: x.index.isin(errors_idx).astype(int))

        if features:
            most_important_idx = []
        elif estimator_name:
            features = []
            model = self._poniard[estimator_name]
            X_train, X_test, y_train, y_test = self._poniard._train_test_split_from_cv()
            model.fit(X_train, y_train)
            scoring = self._poniard._first_scorer(sklearn_scorer=True)
            random_state = self._poniard.random_state
            importances = permutation_importance(
                model,
                X_test,
                y_test,
                n_repeats=10,
                scoring=scoring,
                random_state=random_state,
            )
            sorted_importances_idx = importances.importances_mean.argsort()[::-1]
            if n_features is None:
                n_features = 0.5
            if isinstance(n_features, float):
                assert 0 <= n_features <= 1
                n_features = round(n_features * X.shape[1])
            most_important_idx = sorted_importances_idx[:n_features].tolist()
        else:
            features = X.columns
            most_important_idx = []
        summary = {}
        for i, feature in enumerate(X.columns):
            if (i in most_important_idx or feature in features) and feature != "error":
                current_feature_type = inverted_feature_types[feature]
                if current_feature_type == "numeric":
                    feature_summary = X.groupby("error")[feature].describe()
                else:
                    feature_summary = (
                        X.groupby("error")[feature]
                        .value_counts(normalize=True, dropna=False)
                        .rename("data")
                        .reset_index(feature)
                        .pivot(columns=feature, values="data")
                    )
                summary[feature] = feature_summary
        return summary

    def __repr__(self):
        return non_default_repr(self)
