from __future__ import annotations

from typing import TYPE_CHECKING

import pandas as pd
from sklearn.utils.metaestimators import available_if

from skore._externals._pandas_accessors import DirNamesMixin
from skore._sklearn._base import _BaseAccessor
from skore._sklearn._plot.inspection.coefficients import CoefficientsDisplay
from skore._sklearn._plot.inspection.impurity_decrease import ImpurityDecreaseDisplay
from skore._sklearn._plot.inspection.permutation_importance import (
    PermutationImportanceDisplay,
)
from skore._sklearn.types import DataSource, Metric
from skore._utils._accessor import (
    _check_comparison_report_sub_estimators_have_coef,
    _check_comparison_report_sub_estimators_have_feature_importances,
)

if TYPE_CHECKING:
    from skore import ComparisonReport


class _InspectionAccessor(_BaseAccessor["ComparisonReport"], DirNamesMixin):
    """Accessor for model inspection related operations.

    You can access this accessor using the `inspection` attribute.
    """

    def __init__(self, parent: ComparisonReport) -> None:
        super().__init__(parent)

    @available_if(_check_comparison_report_sub_estimators_have_coef())
    def coefficients(self) -> CoefficientsDisplay:
        """Retrieve the coefficients for each report, including the intercepts.

        Comparison reports with the same features are put under one key and are plotted
        together. When some reports share the same features and others do not, those
        with the same features are plotted together.

        Returns
        -------
        :class:`CoefficientsDisplay`
            The feature importance display containing model coefficients and
            intercept.

        Examples
        --------
        >>> from sklearn.datasets import load_diabetes
        >>> from sklearn.linear_model import Ridge
        >>> from skore import train_test_split
        >>> from skore import ComparisonReport, EstimatorReport
        >>> X, y = load_diabetes(return_X_y=True)
        >>> split_data = train_test_split(X=X, y=y, shuffle=False, as_dict=True)
        >>> report_big_alpha = EstimatorReport(Ridge(alpha=1e3), **split_data)
        >>> report_small_alpha = EstimatorReport(Ridge(alpha=1e-3), **split_data)
        >>> report = ComparisonReport({
        ...     "report small alpha": report_small_alpha,
        ...     "report big alpha": report_big_alpha,
        ... })
        >>> display = report.inspection.coefficients()
        >>> display.frame()
                     estimator     feature  coefficient
        0   report small alpha   Intercept      151.5...
        1   report small alpha  Feature #0      -11.6...
        2   report small alpha  Feature #1     -238.2...
        3   report small alpha  Feature #2      505.5...
        4   report small alpha  Feature #3      298.4...
        5   report small alpha  Feature #4     -408.5...
        6   report small alpha  Feature #5      164.0...
        7   report small alpha  Feature #6      -55.1...
        8   report small alpha  Feature #7      123.0...
        9   report small alpha  Feature #8      633.8...
        10  report small alpha  Feature #9       99.4...
        11    report big alpha   Intercept      151.0...
        12    report big alpha  Feature #0        0.2...
        13    report big alpha  Feature #1        0.0...
        14    report big alpha  Feature #2        0.6...
        15    report big alpha  Feature #3        0.5...
        16    report big alpha  Feature #4        0.2...
        17    report big alpha  Feature #5        0.2...
        18    report big alpha  Feature #6       -0.4...
        19    report big alpha  Feature #7        0.5...
        20    report big alpha  Feature #8        0.6...
        21    report big alpha  Feature #9        0.4...
        >>> display.plot() # shows plot
        """
        return CoefficientsDisplay(
            coefficients=pd.concat(
                [
                    report.inspection.coefficients()
                    .coefficients.copy()
                    .assign(estimator=name)
                    for name, report in self._parent.reports_.items()
                ],
                ignore_index=True,
            ),
            report_type=self._parent._report_type,
        )

    @available_if(_check_comparison_report_sub_estimators_have_feature_importances())
    def impurity_decrease(self) -> ImpurityDecreaseDisplay:
        """Retrieve the Mean Decrease in Impurity (MDI) for each report.

        This method is available for estimators that expose a `feature_importances_`
        attribute. See for example
        :attr:`sklearn.ensemble.GradientBoostingClassifier.inspections_`.

        In particular, note that the MDI is computed at fit time, i.e. using the
        training data.

        Comparison reports with the same features are put under one key and are plotted
        together. When some reports share the same features and others do not, those
        with the same features are plotted together.

        Returns
        -------
        :class:`ImpurityDecreaseDisplay`
            The impurity decrease display containing the feature importances.

        Examples
        --------
        >>> from sklearn.datasets import load_iris
        >>> from sklearn.ensemble import RandomForestClassifier
        >>> from skore import train_test_split
        >>> from skore import ComparisonReport, EstimatorReport
        >>> X, y = load_iris(return_X_y=True, as_frame=True)
        >>> split_data = train_test_split(X=X, y=y, shuffle=False, as_dict=True)
        >>> report_small_trees = EstimatorReport(
        ...     RandomForestClassifier(max_depth=2, random_state=0), **split_data
        ... )
        >>> report_big_trees = EstimatorReport(
        ...     RandomForestClassifier(random_state=0), **split_data
        ... )
        >>> report = ComparisonReport({
        ...     "small trees": report_small_trees,
        ...     "big trees": report_big_trees,
        ... })
        >>> display = report.inspection.impurity_decrease()
        >>> display.frame()
             estimator            feature   importance
        0  small trees  sepal length (cm)       0.1...
        1  small trees   sepal width (cm)       0.0...
        2  small trees  petal length (cm)       0.4...
        3  small trees   petal width (cm)       0.4...
        4    big trees  sepal length (cm)       0.0...
        5    big trees   sepal width (cm)       0.0...
        6    big trees  petal length (cm)       0.4...
        7    big trees   petal width (cm)       0.4...
        >>> display.plot() # shows plot
        """
        return ImpurityDecreaseDisplay(
            importances=pd.concat(
                [
                    report.inspection.impurity_decrease()
                    .importances.copy()
                    .assign(estimator=name)
                    for name, report in self._parent.reports_.items()
                ],
                ignore_index=True,
            ),
            report_type=self._parent._report_type,
        )

    def permutation_importance(
        self,
        *,
        data_source: DataSource = "test",
        at_step: int | str = 0,
        metric: Metric | list[Metric] | dict[str, Metric] | None = None,
        n_repeats: int = 5,
        max_samples: float = 1.0,
        n_jobs: int | None = None,
        seed: int | None = None,
    ) -> PermutationImportanceDisplay:
        """Display the permutation feature importance.

        This computes the permutation importance using sklearn's
        :func:`~sklearn.inspection.permutation_importance` function,
        which consists in permuting the values of one feature and comparing
        the value of `metric` between with and without the permutation, which gives an
        indication on the impact of the feature.

        By default, `seed` is set to `None`, which means the function will
        return a different result at every call. In that case, the results are not
        cached. If you wish to take advantage of skore's caching capabilities, make
        sure you set the `seed` parameter.

        Comparison reports with the same features are put under one key and are plotted
        together. When some reports share the same features and others do not, those
        with the same features are plotted together.

        Parameters
        ----------
        data_source : {"test", "train"}, default="test"
            The data source to use.

            - "test" : use the test set provided when creating the report.
            - "train" : use the train set provided when creating the report.

        at_step : int or str, default=0
            If the estimator is a :class:`~sklearn.pipeline.Pipeline`, at which step of
            the pipeline the importance is computed. If `n`, then the features that
            are evaluated are the ones *right before* the `n`-th step of the pipeline.
            For instance,

            - If 0, compute the importance just before the start of the pipeline (i.e.
              the importance of the raw input features).
            - If -1, compute the importance just before the end of the pipeline (i.e.
              the importance of the fully engineered features, just before the actual
              prediction step).

            If a string, will be searched among the pipeline's `named_steps`.

            Has no effect if the estimator is not a :class:`~sklearn.pipeline.Pipeline`.

        metric : str, callable, scorer, or list of such instances or dict of such \
            instances, default=None
            The metric to pass to :func:`~sklearn.inspection.permutation_importance`.
            The possible values (whether or not in a list) are:

            - if a string, either one of the built-in metrics or a scikit-learn scorer
              name. You can get the possible list of string using
              `report.metrics.help()` or :func:`sklearn.metrics.get_scorer_names` for
              the built-in metrics or the scikit-learn scorers, respectively.
            - if a callable, it should take as arguments `y_true`, `y_pred` as the two
              first arguments. Additional arguments can be passed as keyword arguments
              and will be forwarded with `metric_kwargs`. No favorability indicator can
              be displayed in this case.
            - if the callable API is too restrictive (e.g. need to pass
              same parameter name with different values), you can use scikit-learn
              scorers as provided by :func:`sklearn.metrics.make_scorer`. In this case,
              the metric favorability will only be displayed if it is given explicitly
              via `make_scorer`'s `greater_is_better` parameter.

        n_repeats : int, default=5
            Number of times to permute a feature.

        max_samples : int or float, default=1.0
            The number of samples to draw from `X` to compute feature importance
            in each repeat (without replacement).

            - If int, then draw max_samples samples.
            - If float, then draw max_samples * X.shape[0] samples.
            - If max_samples is equal to 1.0 or X.shape[0], all samples will be used.

            While using this option may provide less accurate importance estimates,
            it keeps the method tractable when evaluating feature importance on
            large datasets. In combination with n_repeats, this allows to control
            the computational speed vs statistical accuracy trade-off of this method.

        n_jobs : int or None, default=None
            Number of jobs to run in parallel. -1 means using all processors.

        seed : int or None, default=None
            The seed used to initialize the random number generator used for the
            permutations.

        Returns
        -------
        :class:`PermutationImportanceDisplay`
            The permutation importance display.

        Examples
        --------
        >>> from sklearn.datasets import make_regression
        >>> from sklearn.linear_model import Ridge
        >>> from sklearn.metrics import make_scorer, r2_score, mean_squared_error
        >>> from skore import train_test_split
        >>> from skore import ComparisonReport, EstimatorReport
        >>> X, y = make_regression(n_features=3, random_state=0)
        >>> split_data = train_test_split(X=X, y=y, shuffle=False, as_dict=True)
        >>> report_big_alpha = EstimatorReport(Ridge(alpha=1e3), **split_data)
        >>> report_small_alpha = EstimatorReport(Ridge(alpha=1e-3), **split_data)
        >>> report = ComparisonReport({
        ...     "small alpha": report_small_alpha,
        ...     "big alpha": report_big_alpha,
        ... })
        >>> report.inspection.permutation_importance(
        ...    n_repeats=2,
        ...    seed=0,
        ... ).frame(aggregate=None)
             estimator data_source metric     feature  repetition     value
        0  small alpha        test     r2  Feature #0           1  0.41...
        1  small alpha        test     r2  Feature #1           1  1.16...
        2  small alpha        test     r2  Feature #2           1  0.01...
        3  small alpha        test     r2  Feature #0           2  0.73...
        4  small alpha        test     r2  Feature #1           2  0.90...
        5  small alpha        test     r2  Feature #2           2  0.01...
        6    big alpha        test     r2  Feature #0           1  0.03...
        7    big alpha        test     r2  Feature #1           1  0.10...
        8    big alpha        test     r2  Feature #2           1  0.00...
        9    big alpha        test     r2  Feature #0           2  0.05...
        10   big alpha        test     r2  Feature #1           2  0.07...
        11   big alpha        test     r2  Feature #2           2 -0.00...

        >>> report.inspection.permutation_importance(
        ...    metric={
        ...        "r2": make_scorer(r2_score),
        ...        "mse": make_scorer(mean_squared_error),
        ...    },
        ...    n_repeats=2,
        ...    seed=0,
        ... ).frame()
              estimator data_source  metric  ...   value_mean    value_std
        0   small alpha        test      r2  ...     0.577888     0.226399
        1   small alpha        test      r2  ...     1.035454     0.180778
        2   small alpha        test      r2  ...     0.016801     0.003577
        3   small alpha        test     mse  ...  3749.565731  1468.964660
        4   small alpha        test     mse  ...  6718.437884  1172.961272
        5   small alpha        test     mse  ...   109.009650    23.205795
        6     big alpha        test      r2  ...     0.044482     0.013717
        7     big alpha        test      r2  ...     0.088916     0.026191
        8     big alpha        test      r2  ...     0.001826     0.008475
        9     big alpha        test     mse  ...   288.618702    88.998809
        10    big alpha        test     mse  ...   576.924696   169.940308
        11    big alpha        test     mse  ...    11.846779    54.992092

        Notes
        -----
        Even if pipeline components output sparse arrays, these will be made dense.
        """
        if seed is not None and not isinstance(seed, int):
            raise ValueError(f"seed must be an integer or None; got {type(seed)}")

        return PermutationImportanceDisplay(
            importances=pd.concat(
                [
                    report.inspection.permutation_importance(
                        data_source=data_source,
                        at_step=at_step,
                        metric=metric,
                        n_repeats=n_repeats,
                        max_samples=max_samples,
                        n_jobs=n_jobs,
                        seed=seed,
                    ).importances.assign(estimator=name)
                    for name, report in self._parent.reports_.items()
                ],
                ignore_index=True,
            ),
            report_type=self._parent._report_type,
        )

    ####################################################################################
    # Methods related to the help tree
    ####################################################################################

    def __repr__(self) -> str:
        """Return a string representation using rich."""
        return self._rich_repr(class_name="skore.EstimatorReport.inspection")
