from collections import defaultdict

import matplotlib.pyplot as plt
import numpy as np
from sklearn.metrics import average_precision_score, precision_recall_curve
from sklearn.preprocessing import LabelBinarizer

from skore.sklearn._plot.utils import (
    HelpDisplayMixin,
    _ClassifierCurveDisplayMixin,
    _despine_matplotlib_axis,
    _validate_style_kwargs,
    sample_mpl_colormap,
)


class PrecisionRecallCurveDisplay(HelpDisplayMixin, _ClassifierCurveDisplayMixin):
    """Precision Recall visualization.

    An instance of this class is should created by
    `EstimatorReport.metrics.plot.precision_recall()`. You should not create an
    instance of this class directly.


    Parameters
    ----------
    precision : dict of list of ndarray
        Precision values. The structure is:

        - for binary classification:
            - the key is the positive label.
            - the value is a list of `ndarray`, each `ndarray` being the precision.
        - for multiclass classification:
            - the key is the class of interest in an OvR fashion.
            - the value is a list of `ndarray`, each `ndarray` being the precision.

    recall : dict of list of ndarray
        Recall values. The structure is:

        - for binary classification:
            - the key is the positive label.
            - the value is a list of `ndarray`, each `ndarray` being the recall.
        - for multiclass classification:
            - the key is the class of interest in an OvR fashion.
            - the value is a list of `ndarray`, each `ndarray` being the recall.

    average_precision : dict of list of float
        Average precision. The structure is:

        - for binary classification:
            - the key is the positive label.
            - the value is a list of `float`, each `float` being the average
              precision.
        - for multiclass classification:
            - the key is the class of interest in an OvR fashion.
            - the value is a list of `float`, each `float` being the average
              precision.

    estimator_name : str
        Name of the estimator.

    pos_label : int, float, bool or str, default=None
        The class considered as the positive class. If None, the class will not
        be shown in the legend.

    data_source : {"train", "test", "X_y"}, default=None
        The data source used to compute the precision recall curve.

    Attributes
    ----------
    ax_ : matplotlib Axes
        Axes with precision recall curve.

    figure_ : matplotlib Figure
        Figure containing the curve.

    lines_ : list of matplotlib Artist
        Precision recall curve.

    chance_levels_ : matplotlib Artist or None
        The chance level line. It is `None` if the chance level is not plotted.

    Examples
    --------
    >>> from sklearn.datasets import load_breast_cancer
    >>> from sklearn.linear_model import LogisticRegression
    >>> from sklearn.model_selection import train_test_split
    >>> from skore import EstimatorReport
    >>> X_train, X_test, y_train, y_test = train_test_split(
    ...     *load_breast_cancer(return_X_y=True), random_state=0
    ... )
    >>> classifier = LogisticRegression(max_iter=10_000)
    >>> reporter = EstimatorReport(
    ...     classifier,
    ...     X_train=X_train,
    ...     y_train=y_train,
    ...     X_test=X_test,
    ...     y_test=y_test,
    ... )
    >>> display = reporter.metrics.plot.precision_recall()
    >>> display.plot(pr_curve_kwargs={"color": "tab:red"})
    """

    def __init__(
        self,
        precision,
        recall,
        *,
        average_precision,
        estimator_name,
        pos_label=None,
        data_source=None,
    ):
        self.precision = precision
        self.recall = recall
        self.average_precision = average_precision
        self.estimator_name = estimator_name
        self.pos_label = pos_label
        self.data_source = data_source

    def plot(
        self,
        ax=None,
        *,
        estimator_name=None,
        pr_curve_kwargs=None,
        despine=True,
    ):
        """Plot visualization.

        Extra keyword arguments will be passed to matplotlib's `plot`.

        Parameters
        ----------
        ax : Matplotlib Axes, default=None
            Axes object to plot on. If `None`, a new figure and axes is
            created.

        estimator_name : str, default=None
            Name of the estimator used to plot the precision-recall curve. If
            `None`, we use the inferred name from the estimator.

        pr_curve_kwargs : dict or list of dict, default=None
            Keyword arguments to be passed to matplotlib's `plot` for rendering
            the precision-recall curve(s).

        despine : bool, default=True
            Whether to remove the top and right spines from the plot.

        Returns
        -------
        display : PrecisionRecallCurveDisplay
            Object that stores computed values.

        Notes
        -----
        The average precision (cf. :func:`~sklearn.metrics.average_precision_score`)
        in scikit-learn is computed without any interpolation. To be consistent
        with this metric, the precision-recall curve is plotted without any
        interpolation as well (step-wise style).

        You can change this style by passing the keyword argument
        `drawstyle="default"`. However, the curve will not be strictly
        consistent with the reported average precision.

        Examples
        --------
        >>> from sklearn.datasets import load_breast_cancer
        >>> from sklearn.linear_model import LogisticRegression
        >>> from sklearn.model_selection import train_test_split
        >>> from skore import EstimatorReport
        >>> X_train, X_test, y_train, y_test = train_test_split(
        ...     *load_breast_cancer(return_X_y=True), random_state=0
        ... )
        >>> classifier = LogisticRegression(max_iter=10_000)
        >>> reporter = EstimatorReport(
        ...     classifier,
        ...     X_train=X_train,
        ...     y_train=y_train,
        ...     X_test=X_test,
        ...     y_test=y_test,
        ... )
        >>> display = reporter.metrics.plot.precision_recall()
        >>> display.plot(pr_curve_kwargs={"color": "tab:red"})
        """
        self.ax_, self.figure_, estimator_name = self._validate_plot_params(
            ax=ax, estimator_name=estimator_name
        )

        self.lines_ = []
        self.chance_levels_ = []
        if len(self.precision) == 1:  # binary-classification
            if len(self.precision[self.pos_label]) == 1:  # single-split
                if pr_curve_kwargs is None:
                    pr_curve_kwargs = {}
                elif isinstance(pr_curve_kwargs, list):
                    if len(pr_curve_kwargs) > 1:
                        raise ValueError(
                            "You intend to plot a single precision-recall curve and "
                            "provide multiple precision-recall curve keyword "
                            "arguments. Provide a single dictionary or a list with "
                            "a single dictionary."
                        )
                    pr_curve_kwargs = pr_curve_kwargs[0]

                precision = self.precision[self.pos_label][0]
                recall = self.recall[self.pos_label][0]
                average_precision = self.average_precision[self.pos_label][0]

                default_line_kwargs = {"drawstyle": "steps-post"}
                if self.data_source in ("train", "test"):
                    default_line_kwargs["label"] = (
                        f"{self.data_source.title()} set "
                        f"(AP = {average_precision:0.2f})"
                    )
                else:  # data_source in (None, "X_y")
                    default_line_kwargs["label"] = f"AP = {average_precision:0.2f}"

                line_kwargs = _validate_style_kwargs(
                    default_line_kwargs, pr_curve_kwargs
                )

                (line_,) = self.ax_.plot(recall, precision, **line_kwargs)
                self.lines_.append(line_)
            else:  # cross-validation
                if pr_curve_kwargs is None:
                    pr_curve_kwargs = [{}] * len(self.precision[self.pos_label])
                elif isinstance(pr_curve_kwargs, list):
                    if len(pr_curve_kwargs) != len(self.precision[self.pos_label]):
                        raise ValueError(
                            "You intend to plot multiple precision-recall curves. We "
                            "expect `pr_curve_kwargs` to be a list of dictionaries "
                            "with the same length as the number of precision-recall "
                            "curves. Got "
                            f"{len(pr_curve_kwargs)} instead of "
                            f"{len(self.precision)}."
                        )
                else:
                    raise ValueError(
                        "You intend to plot multiple precision-recall curves. We "
                        "expect `pr_curve_kwargs` to be a list of dictionaries of "
                        f"{len(self.precision)} elements. Got {pr_curve_kwargs!r} "
                        "instead."
                    )

                for split_idx in range(len(self.precision[self.pos_label])):
                    precision = self.precision[self.pos_label][split_idx]
                    recall = self.recall[self.pos_label][split_idx]
                    average_precision = self.average_precision[self.pos_label][
                        split_idx
                    ]

                    default_line_kwargs = {
                        "drawstyle": "steps-post",
                        "label": (
                            f"{self.data_source.title()} set - fold #{split_idx + 1} "
                            f"(AP = {average_precision:0.2f})"
                        ),
                    }
                    line_kwargs = _validate_style_kwargs(
                        default_line_kwargs, pr_curve_kwargs[split_idx]
                    )

                    (line_,) = self.ax_.plot(recall, precision, **line_kwargs)
                    self.lines_.append(line_)

            info_pos_label = (
                f"\n(Positive label: {self.pos_label})"
                if self.pos_label is not None
                else ""
            )
        else:  # multiclass-classification
            info_pos_label = None  # irrelevant for multiclass
            class_colors = sample_mpl_colormap(
                plt.cm.tab10, 10 if len(self.precision) < 10 else len(self.precision)
            )
            if pr_curve_kwargs is None:
                pr_curve_kwargs = [{}] * len(self.precision)
            elif isinstance(pr_curve_kwargs, list):
                if len(pr_curve_kwargs) != len(self.precision):
                    raise ValueError(
                        "You intend to plot multiple precision-recall curves. We "
                        "expect `pr_curve_kwargs` to be a list of dictionaries with "
                        "the same length as the number of precision-recall curves. "
                        "Got "
                        f"{len(pr_curve_kwargs)} instead of "
                        f"{len(self.precision)}."
                    )
            else:
                raise ValueError(
                    "You intend to plot multiple precision-recall curves. We expect "
                    "`pr_curve_kwargs` to be a list of dictionaries of "
                    f"{len(self.precision)} elements. Got {pr_curve_kwargs!r} instead."
                )

            for class_idx, class_ in enumerate(self.precision):
                precision_class = self.precision[class_]
                recall_class = self.recall[class_]
                average_precision_class = self.average_precision[class_]
                pr_curve_kwargs_class = pr_curve_kwargs[class_idx]

                if len(precision_class) == 1:  # single-split
                    precision = precision_class[0]
                    recall = recall_class[0]
                    average_precision = average_precision_class[0]

                    default_line_kwargs = {
                        "drawstyle": "steps-post",
                        "color": class_colors[class_idx],
                    }
                    if self.data_source in ("train", "test"):
                        default_line_kwargs["label"] = (
                            f"{str(class_).title()} - {self.data_source} set "
                            f"(AP = {average_precision:0.2f})"
                        )
                    else:  # data_source in (None, "X_y")
                        default_line_kwargs["label"] = (
                            f"{str(class_).title()} AP = {average_precision:0.2f}"
                        )

                    line_kwargs = _validate_style_kwargs(
                        default_line_kwargs, pr_curve_kwargs_class
                    )

                    (line_,) = self.ax_.plot(recall, precision, **line_kwargs)
                    self.lines_.append(line_)
                else:  # cross-validation
                    for split_idx in range(len(precision_class)):
                        precision = precision_class[split_idx]
                        recall = recall_class[split_idx]
                        average_precision = average_precision_class[split_idx]

                        default_line_kwargs = {
                            "color": class_colors[class_idx],
                            "alpha": 0.3,
                        }
                        if split_idx == 0:
                            default_line_kwargs["label"] = (
                                f"{str(class_).title()} - {self.data_source} set"
                                f" (AP = {np.mean(average_precision_class):0.2f} +/- "
                                f"{np.std(average_precision_class):0.2f})"
                            )
                        else:
                            default_line_kwargs["label"] = None

                        line_kwargs = _validate_style_kwargs(default_line_kwargs, {})

                        (line_,) = self.ax_.plot(recall, precision, **line_kwargs)
                        self.lines_.append(line_)

        xlabel = "Recall"
        ylabel = "Precision"
        if info_pos_label:
            xlabel += info_pos_label
            ylabel += info_pos_label

        self.ax_.set(
            xlabel=xlabel,
            xlim=(-0.01, 1.01),
            ylabel=ylabel,
            ylim=(-0.01, 1.01),
            aspect="equal",
        )

        if despine:
            _despine_matplotlib_axis(self.ax_)

        self.ax_.legend(loc="lower left", title=estimator_name)

    @classmethod
    def _from_predictions(
        cls,
        y_true,
        y_pred,
        *,
        estimator,
        estimator_name,
        ml_task,
        data_source=None,
        pos_label=None,
        drop_intermediate=False,
        ax=None,
        pr_curve_kwargs=None,
        despine=True,
    ):
        """Plot precision-recall curve given binary class predictions.

        Parameters
        ----------
        y_true : array-like of shape (n_samples,)
            True binary labels.

        y_pred : array-like of shape (n_samples,)
            Target scores, can either be probability estimates of the positive class,
            confidence values, or non-thresholded measure of decisions (as returned by
            “decision_function” on some classifiers).

        estimator : estimator instance
            The estimator from which `y_pred` is obtained.

        estimator_name : str
            Name of the estimator used to plot the precision-recall curve.

        ml_task : {"binary-classification", "multiclass-classification"}
            The machine learning task.

        data_source : {"train", "test", "X_y"}, default=None
            The data source used to compute the ROC curve.

        pos_label : int, float, bool or str, default=None
            The class considered as the positive class when computing the
            precision and recall metrics.

        drop_intermediate : bool, default=False
            Whether to drop some suboptimal thresholds which would not appear
            on a plotted precision-recall curve. This is useful in order to
            create lighter precision-recall curves.

        ax : matplotlib axes, default=None
            Axes object to plot on. If `None`, a new figure and axes is created.

        pr_curve_kwargs : dict or list of dict, default=None
            Keyword arguments to be passed to matplotlib's `plot` for rendering
            the precision-recall curve(s).

        despine : bool, default=True
            Whether to remove the top and right spines from the plot.

        **kwargs : dict
            Keyword arguments to be passed to matplotlib's `plot`.

        Returns
        -------
        display : :class:`~sklearn.metrics.PrecisionRecallDisplay`
        """
        pos_label_validated = cls._validate_from_predictions_params(
            y_true, y_pred, ml_task=ml_task, pos_label=pos_label
        )

        if ml_task == "binary-classification":
            precision, recall = defaultdict(list), defaultdict(list)
            average_precision = defaultdict(list)
            for y_true_i, y_pred_i in zip(y_true, y_pred):
                precision_i, recall_i, _ = precision_recall_curve(
                    y_true_i,
                    y_pred_i,
                    pos_label=pos_label_validated,
                    drop_intermediate=drop_intermediate,
                )
                average_precision_i = average_precision_score(
                    y_true_i, y_pred_i, pos_label=pos_label_validated
                )

                precision[pos_label_validated].append(precision_i)
                recall[pos_label_validated].append(recall_i)
                average_precision[pos_label_validated].append(average_precision_i)
        else:  # multiclass-classification
            precision, recall = defaultdict(list), defaultdict(list)
            average_precision = defaultdict(list)
            for y_true_i, y_pred_i in zip(y_true, y_pred):
                label_binarizer = LabelBinarizer().fit(estimator.classes_)
                y_true_onehot_i = label_binarizer.transform(y_true_i)
                for class_idx, class_ in enumerate(estimator.classes_):
                    precision_class_i, recall_class_i, _ = precision_recall_curve(
                        y_true_onehot_i[:, class_idx],
                        y_pred_i[:, class_idx],
                        pos_label=None,
                        drop_intermediate=drop_intermediate,
                    )
                    average_precision_class_i = average_precision_score(
                        y_true_onehot_i[:, class_idx], y_pred_i[:, class_idx]
                    )

                    precision[class_].append(precision_class_i)
                    recall[class_].append(recall_class_i)
                    average_precision[class_].append(average_precision_class_i)

        viz = cls(
            precision=precision,
            recall=recall,
            average_precision=average_precision,
            estimator_name=estimator_name,
            pos_label=pos_label_validated,
            data_source=data_source,
        )

        viz.plot(
            ax=ax,
            estimator_name=estimator_name,
            pr_curve_kwargs=pr_curve_kwargs,
            despine=despine,
        )

        return viz
