import contextlib
import copy
from typing import Any, Literal

import numpy as np
import pandas as pd
from matplotlib.legend import Legend
from sklearn.base import BaseEstimator, ClassifierMixin
from sklearn.metrics import r2_score

from skore import configuration
from skore._sklearn._base import _BaseAccessor, _BaseReport
from skore._sklearn._plot.metrics.precision_recall_curve import (
    PrecisionRecallCurveDisplay,
)
from skore._sklearn._plot.metrics.roc_curve import RocCurveDisplay


def check_precision_recall_curve_display_data(display: PrecisionRecallCurveDisplay):
    """Check the structure of the display's internal data."""
    assert list(display.precision_recall.columns) == [
        "estimator",
        "data_source",
        "split",
        "label",
        "threshold",
        "precision",
        "recall",
    ]
    assert list(display.average_precision.columns) == [
        "estimator",
        "data_source",
        "split",
        "label",
        "average_precision",
    ]


@contextlib.contextmanager
def check_cache_changed(value):
    """Assert that `value` has changed during context execution."""
    initial_value = copy.copy(value)
    yield
    assert value != initial_value


@contextlib.contextmanager
def check_cache_unchanged(value):
    """Assert that `value` has not changed during context execution."""
    initial_value = copy.copy(value)
    yield
    assert value == initial_value


class MockEstimator(ClassifierMixin, BaseEstimator):
    def __init__(self, *, error, n_call=0, fail_after_n_clone=3):
        self.error = error
        self.n_call = n_call
        self.fail_after_n_clone = fail_after_n_clone

    def fit(self, X, y):
        if self.n_call > self.fail_after_n_clone:
            raise self.error
        self.classes_ = np.unique(y)
        return self

    def __sklearn_clone__(self):
        self.n_call += 1
        return self

    def predict(self, X):
        return np.ones(X.shape[0])


class MockReport(_BaseReport):
    """Minimal report for testing.

    Attributes
    ----------
    no_private : object
        Public attribute for tests.
    """

    _ACCESSOR_CONFIG: dict[str, dict[str, str]] = {}

    def __init__(
        self,
        estimator,
        X_train=None,
        y_train=None,
        X_test=None,
        y_test=None,
    ):
        self._estimator = estimator
        self._X_train = X_train
        self._y_train = y_train
        self._X_test = X_test
        self._y_test = y_test
        self.no_private = "no_private"
        self.attr_without_description = "attr_without_description"

    @property
    def estimator_(self):
        return self._estimator

    @property
    def X_train(self):
        return self._X_train

    @property
    def y_train(self):
        return self._y_train

    @property
    def X_test(self):
        return self._X_test

    @property
    def y_test(self):
        return self._y_test

    def _get_help_title(self) -> str:
        return "Mock report"


class MockAccessor(_BaseAccessor):
    """Minimal accessor for testing."""

    _accessor_name = "mock_accessor"  # registering an accessor will set this attribute

    def __init__(self, parent):
        super().__init__(parent)

    def _get_help_tree_title(self) -> str:
        return "Mock accessor"


class MockDisplay:
    """Minimal display for testing, following the Display protocol."""

    def plot(self, **kwargs: Any) -> None:
        pass

    def set_style(
        self, *, policy: Literal["override", "update"] = "update", **kwargs: Any
    ) -> None:
        pass

    def frame(self, **kwargs: Any) -> pd.DataFrame:
        return pd.DataFrame()

    def help(self) -> None:
        pass


def check_roc_curve_display_data(display: RocCurveDisplay):
    """Check the structure of the display's internal data."""
    assert list(display.roc_curve.columns) == [
        "estimator",
        "data_source",
        "split",
        "label",
        "threshold",
        "fpr",
        "tpr",
    ]
    assert list(display.roc_auc.columns) == [
        "estimator",
        "data_source",
        "split",
        "label",
        "roc_auc",
    ]


def check_legend_position(ax, *, loc: str, position: Literal["inside", "outside"]):
    """Check the position of the legend in the axes."""
    legend = ax.get_legend()
    assert legend._loc == Legend.codes[loc]
    bbox = legend.get_window_extent().transformed(ax.transAxes.inverted())
    tol = 1e-6
    if position == "inside":
        assert -tol <= bbox.x0 <= 1 + tol
    else:
        assert bbox.x0 >= 1 - tol


def check_frame_structure(df, expected_index, expected_data_columns):
    """Check the structure of ROC, precision-recall, and prediction error DataFrames.

    Parameters
    ----------
    df : DataFrame
        The DataFrame to check.
    expected_index : list of str
        The expected index column names (e.g., `estimator_name`, `split`,
        `label`).
        These columns should be of categorical type.
    expected_data_columns : list of str
        The expected data column names (e.g., `threshold`, `fpr`, `tpr`,
        `precision`, `recall`).
        These columns should be of `np.float64` type.
    """
    assert isinstance(df, pd.DataFrame)
    assert sorted(df.columns.tolist()) == sorted(expected_index + expected_data_columns)

    for col in df.columns:
        if col in expected_index:
            assert df[col].dtype.name == "category"
        else:
            assert col in expected_data_columns
            assert df[col].dtype == np.float64


def custom_r2_score(estimator, X, y):
    """Compute the R^2 score.

    This is a custom callable scorer for testing permutation importance with callable
    metrics.

    This function has the signature required by sklearn's permutation_importance:
    (estimator, X, y) -> score.

    Parameters
    ----------
    estimator : object
        The fitted estimator.
    X : array-like
        The input data.
    y : array-like
        The target values.

    Returns
    -------
    float
        The R2 score.
    """
    y_pred = estimator.predict(X)
    return r2_score(y, y_pred)


def _change_configuration_for_testing():
    """
    Change configuration for testing purposes, especially on ``joblib.Parallel``.

    This function should exist in a Python module rather than in tests, otherwise joblib
    will not be able to pickle it.
    """
    show_progress_copy = configuration.show_progress
    plot_backend_copy = configuration.plot_backend

    configuration.show_progress = "show_progress_thread"
    configuration.plot_backend = "plot_backend_thread"

    return (
        (show_progress_copy, plot_backend_copy),
        (configuration.show_progress, configuration.plot_backend),
    )
