from __future__ import annotations

import uuid
from collections import Counter
from collections.abc import Iterable
from typing import TYPE_CHECKING, Literal, cast

import joblib
import numpy as np
import pandas as pd
from numpy.typing import ArrayLike

from skore._externals._pandas_accessors import DirNamesMixin
from skore._sklearn._base import _BaseReport
from skore._sklearn._cross_validation.report import CrossValidationReport
from skore._sklearn._estimator.report import EstimatorReport
from skore._sklearn.types import PositiveLabel
from skore._utils._progress_bar import track
from skore._utils.repr.data import get_documentation_url
from skore._utils.repr.html_repr import render_template

if TYPE_CHECKING:
    from skore._sklearn._comparison.inspection_accessor import (
        _InspectionAccessor,
    )
    from skore._sklearn._comparison.metrics_accessor import _MetricsAccessor

    ComparisonReportType = Literal[
        "comparison-estimator",
        "comparison-cross-validation",
    ]


class ComparisonReport(_BaseReport, DirNamesMixin):
    """Report for comparing reports.

    This object can be used to compare several :class:`skore.EstimatorReport` instances,
    or several :class:`~skore.CrossValidationReport` instances.

    Refer to the :ref:`comparison_report` section of the user guide for more details.

    .. caution::
       Reports passed to `ComparisonReport` are not copied. If you pass
       a report to `ComparisonReport`, and then modify the report outside later, it
       will affect the report stored inside the `ComparisonReport` as well, which
       can lead to inconsistent results. For this reason, modifying reports after
       creation is strongly discouraged.

    Parameters
    ----------
    reports : list of reports or dict
        Reports to compare. If a dict, keys will be used to label the estimators;
        if a list, the labels are computed from the estimator class names.
        Expects at least two reports to compare, with the same test target.

    n_jobs : int, default=None
        Number of jobs to run in parallel. Training the estimators and computing
        the scores are parallelized.
        When accessing some methods of the `ComparisonReport`, the `n_jobs`
        parameter is used to parallelize the computation.
        ``None`` means 1 unless in a :obj:`joblib.parallel_backend` context.
        ``-1`` means using all processors.

    Attributes
    ----------
    reports_ : dict mapping names to reports
        The compared reports.

    See Also
    --------
    skore.EstimatorReport
        Report for a fitted estimator.

    skore.CrossValidationReport
        Report for the cross-validation of an estimator.

    Examples
    --------
    >>> from sklearn.datasets import make_classification
    >>> from skore import train_test_split
    >>> from sklearn.linear_model import LogisticRegression
    >>> from skore import ComparisonReport, EstimatorReport
    >>> X, y = make_classification(random_state=42)
    >>> split_data = train_test_split(X=X, y=y, random_state=42, as_dict=True)
    >>> estimator_1 = LogisticRegression()
    >>> estimator_report_1 = EstimatorReport(estimator_1, **split_data)
    >>> estimator_2 = LogisticRegression(C=2)  # Different regularization
    >>> estimator_report_2 = EstimatorReport(estimator_2, **split_data)
    >>> report = ComparisonReport([estimator_report_1, estimator_report_2])
    >>> report.reports_
    {'LogisticRegression_1': ..., 'LogisticRegression_2': ...}
    >>> report = ComparisonReport(
    ...     {"model1": estimator_report_1, "model2": estimator_report_2}
    ... )
    >>> report.reports_
    {'model1': ..., 'model2': ...}

    >>> from sklearn.datasets import make_classification
    >>> from sklearn.linear_model import LogisticRegression
    >>> from skore import ComparisonReport, CrossValidationReport
    >>> X, y = make_classification(random_state=42)
    >>> estimator_1 = LogisticRegression()
    >>> estimator_2 = LogisticRegression(C=2)  # Different regularization
    >>> report_1 = CrossValidationReport(estimator_1, X, y)
    >>> report_2 = CrossValidationReport(estimator_2, X, y)
    >>> report = ComparisonReport([report_1, report_2])
    >>> report = ComparisonReport({"model1": report_1, "model2": report_2})
    """

    _ACCESSOR_CONFIG: dict[str, dict[str, str]] = {
        "metrics": {"name": "metrics"},
        "inspection": {"name": "inspection"},
    }
    metrics: _MetricsAccessor
    inspection: _InspectionAccessor

    _report_type: ComparisonReportType

    @staticmethod
    def _validate_reports(
        reports: (
            list[EstimatorReport]
            | dict[str, EstimatorReport]
            | list[CrossValidationReport]
            | dict[str, CrossValidationReport]
        ),
    ) -> tuple[
        dict[str, EstimatorReport] | dict[str, CrossValidationReport],
        ComparisonReportType,
        PositiveLabel,
    ]:
        """Validate that reports are in the right format for comparison.

        Parameters
        ----------
        reports : list of reports or dict
            The reports to be validated.

        Returns
        -------
        dict
            The validated reports.
        {"comparison-estimator", "comparison-cross-validation"}
            The report type for the comparison.
        int, float, bool, str or None
            The positive label used in the different reports.
        """
        if not isinstance(reports, Iterable):
            raise TypeError(
                f"Expected reports to be a list or dict; got {type(reports)}"
            )

        if len(reports) < 2:
            raise ValueError(
                f"Expected at least 2 reports to compare; got {len(reports)}"
            )

        if isinstance(reports, list):
            report_names = None
            reports_list = reports
        else:  # dict
            report_names = list(reports.keys())
            for key in report_names:
                if not isinstance(key, str):
                    raise TypeError(
                        f"Expected all report names to be strings; got {type(key)}"
                    )
            reports_list = cast(
                list[EstimatorReport] | list[CrossValidationReport],
                list(reports.values()),
            )

        report_type: ComparisonReportType
        if all(isinstance(report, EstimatorReport) for report in reports_list):
            reports_list = cast(list[EstimatorReport], reports_list)
            report_type = "comparison-estimator"

            test_dataset_hashes = {
                joblib.hash(report.y_test)
                for report in reports_list
                if report.y_test is not None
            }
            if len(test_dataset_hashes) > 1:
                raise ValueError(
                    "Expected all estimators to share the same test targets."
                )

        elif all(isinstance(report, CrossValidationReport) for report in reports_list):
            reports_list = cast(list[CrossValidationReport], reports_list)
            report_type = "comparison-cross-validation"
        else:
            raise TypeError(
                f"Expected list or dict of {EstimatorReport.__name__} "
                f"or list of dict of {CrossValidationReport.__name__}"
            )

        if len({id(report) for report in reports_list}) < len(reports_list):
            raise ValueError("Expected reports to be distinct objects")

        ml_tasks = {report._ml_task for report in reports_list}
        if len(ml_tasks) > 1:
            raise ValueError(
                f"Expected all estimators to have the same ML usecase; got {ml_tasks}"
            )

        if ml_tasks == {"binary-classification"}:
            pos_labels = {report.pos_label for report in reports_list}
            if len(pos_labels) > 1:
                raise ValueError(
                    "Expected all estimators to have the same positive label. "
                    f"Got {pos_labels}."
                )
            pos_label = pos_labels.pop()
        else:
            pos_label = None

        if report_names is None:
            deduped_report_names = _deduplicate_report_names(
                [report.estimator_name_ for report in reports_list]
            )
        else:
            deduped_report_names = report_names

        reports_dict = cast(
            dict[str, EstimatorReport] | dict[str, CrossValidationReport],
            dict(zip(deduped_report_names, reports_list, strict=True)),
        )

        return reports_dict, report_type, pos_label

    def __init__(
        self,
        reports: (
            list[EstimatorReport]
            | dict[str, EstimatorReport]
            | list[CrossValidationReport]
            | dict[str, CrossValidationReport]
        ),
        *,
        n_jobs: int | None = None,
    ) -> None:
        """
        ComparisonReport instance initializer.

        Notes
        -----
        We check that the estimator reports can be compared:
        - all reports are estimator reports,
        - all estimators are in the same ML use case,
        - all estimators have non-empty X_test and y_test,
        - all estimators have the same X_test and y_test.
        """
        super().__init__()
        self.reports_, self._report_type, self._pos_label = (
            ComparisonReport._validate_reports(reports)
        )

        self.n_jobs = n_jobs
        self._ml_task = next(iter(self.reports_.values()))._ml_task  # type: ignore

    def clear_cache(self) -> None:
        """Clear the cache.

        Examples
        --------
        >>> from sklearn.datasets import make_classification
        >>> from sklearn.linear_model import LogisticRegression
        >>> from skore import train_test_split
        >>> from skore import ComparisonReport, EstimatorReport
        >>> X, y = make_classification(random_state=42)
        >>> split_data = train_test_split(X=X, y=y, random_state=42, as_dict=True)
        >>> estimator_1 = LogisticRegression()
        >>> estimator_report_1 = EstimatorReport(estimator_1, **split_data)
        >>> estimator_2 = LogisticRegression(C=2)  # Different regularization
        >>> estimator_report_2 = EstimatorReport(estimator_2, **split_data)
        >>> report = ComparisonReport([estimator_report_1, estimator_report_2])
        >>> report.cache_predictions()
        >>> report.clear_cache()
        """
        for report in self.reports_.values():
            report.clear_cache()

    def cache_predictions(
        self,
        response_methods: Literal[
            "auto", "predict", "predict_proba", "decision_function"
        ] = "auto",
        n_jobs: int | None = None,
    ) -> None:
        """Cache the predictions for sub-estimators reports.

        Parameters
        ----------
        response_methods : {"auto", "predict", "predict_proba", "decision_function"},\
                default="auto"
            The methods to use to compute the predictions.

        n_jobs : int, default=None
            The number of jobs to run in parallel. If `None`, we use the `n_jobs`
            parameter when initializing the report.

        Examples
        --------
        >>> from sklearn.datasets import make_classification
        >>> from sklearn.linear_model import LogisticRegression
        >>> from skore import train_test_split
        >>> from skore import ComparisonReport, EstimatorReport
        >>> X, y = make_classification(random_state=42)
        >>> split_data = train_test_split(X=X, y=y, random_state=42, as_dict=True)
        >>> estimator_1 = LogisticRegression()
        >>> estimator_report_1 = EstimatorReport(estimator_1, **split_data)
        >>> estimator_2 = LogisticRegression(C=2)  # Different regularization
        >>> estimator_report_2 = EstimatorReport(estimator_2, **split_data)
        >>> report = ComparisonReport([estimator_report_1, estimator_report_2])
        >>> report.cache_predictions()
        """
        if n_jobs is None:
            n_jobs = self.n_jobs

        for report in track(
            self.reports_.values(),
            description="Estimator predictions",
            total=len(self.reports_),
        ):
            report.cache_predictions(response_methods=response_methods, n_jobs=n_jobs)

    def get_predictions(
        self,
        *,
        data_source: Literal["train", "test"],
        response_method: Literal[
            "predict", "predict_proba", "decision_function"
        ] = "predict",
    ) -> list[ArrayLike] | list[list[ArrayLike]]:
        """Get predictions from the underlying reports.

        This method has the advantage to reload from the cache if the predictions
        were already computed in a previous call.

        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.

        response_method : {"predict", "predict_proba", "decision_function"}, \
                default="predict"
            The response method to use to get the predictions.

        Returns
        -------
        list of np.ndarray of shape (n_samples,) or (n_samples, n_classes) or list of \
                such lists
            The predictions for each :class:`~skore.EstimatorReport` or
            :class:`~skore.CrossValidationReport`.

        Raises
        ------
        ValueError
            If the data source is invalid.

        Examples
        --------
        >>> from sklearn.datasets import make_classification
        >>> from skore import train_test_split
        >>> from sklearn.linear_model import LogisticRegression
        >>> from skore import ComparisonReport, EstimatorReport
        >>> X, y = make_classification(random_state=42)
        >>> split_data = train_test_split(X=X, y=y, random_state=42, as_dict=True)
        >>> estimator_1 = LogisticRegression()
        >>> estimator_report_1 = EstimatorReport(estimator_1, **split_data)
        >>> estimator_2 = LogisticRegression(C=2)  # Different regularization
        >>> estimator_report_2 = EstimatorReport(estimator_2, **split_data)
        >>> report = ComparisonReport([estimator_report_1, estimator_report_2])
        >>> report.cache_predictions()
        >>> predictions = report.get_predictions(data_source="test")
        >>> print([split_predictions.shape for split_predictions in predictions])
        [(25,), (25,)]
        """
        return [  # type: ignore
            report.get_predictions(
                data_source=data_source,
                response_method=response_method,
            )
            for report in self.reports_.values()
        ]

    def create_estimator_report(
        self,
        *,
        report_key: str,
        X_test: ArrayLike,
        y_test: ArrayLike,
    ) -> EstimatorReport:
        """Create an estimator report from one of the reports in the comparison.

        This method creates a new :class:`~skore.EstimatorReport` with the same
        estimator and the same data as the chosen report. It is useful to evaluate and
        deploy a model that was deemed optimal during the comparison. Provide a held out
        test set to properly evaluate the performance of the model.

        Parameters
        ----------
        report_key : str
            The key associated with the estimator to create a report for, as stored in
            the :attr:`reports_` attribute of the :class:`~skore.ComparisonReport`.
            List the available keys with `reports_.keys()`.

        X_test : {array-like, sparse matrix} of shape (n_samples, n_features)
            Testing data. It should have the same structure as the training data.

        y_test : array-like of shape (n_samples,) or (n_samples, n_outputs) or None
            Testing target.

        Examples
        --------
        >>> from sklearn.datasets import make_classification
        >>> from sklearn.ensemble import RandomForestClassifier
        >>> from sklearn.linear_model import LogisticRegression
        >>> from skore import train_test_split
        >>> from skore import ComparisonReport, CrossValidationReport
        >>> X, y = make_classification(random_state=42)
        >>> X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=42)
        >>> linear_report = CrossValidationReport(
        ...     LogisticRegression(random_state=42), X_train, y_train
        ... )
        >>> forest_report = CrossValidationReport(
        ...     RandomForestClassifier(random_state=42), X_train, y_train
        ... )
        >>> comparison_report = ComparisonReport([linear_report, forest_report])
        >>> summary = comparison_report.metrics.summarize().frame()

        >>> # Notice that e.g. the RandomForestClassifier performs best
        >>> final_report = comparison_report.create_estimator_report(
        ...     report_key="RandomForestClassifier", X_test=X_test, y_test=y_test
        ... )
        >>> final_report.metrics.summarize().frame()

        Returns
        -------
        report : :class:`~skore.EstimatorReport`
            The estimator report.
        """
        if report_key not in self.reports_:
            raise ValueError(
                f"Estimator with key {report_key} not found in the comparison report. "
                f"Available keys are: {list(self.reports_.keys())}."
            )

        if isinstance(self.reports_[report_key], CrossValidationReport):
            return cast(
                CrossValidationReport, self.reports_[report_key]
            ).create_estimator_report(X_test=X_test, y_test=y_test)

        estimator_report = cast(EstimatorReport, self.reports_[report_key])
        X_concat = (
            pd.concat([estimator_report._X_train, estimator_report._X_test])
            if isinstance(estimator_report._X_train, pd.DataFrame)
            else np.concatenate([estimator_report._X_train, estimator_report._X_test])
        )
        y_concat = (
            pd.concat([estimator_report._y_train, estimator_report._y_test])
            if isinstance(estimator_report._y_train, (pd.DataFrame, pd.Series))
            else np.concatenate([estimator_report._y_train, estimator_report._y_test])
        )
        report = EstimatorReport(
            estimator_report.estimator,
            X_train=X_concat,
            y_train=y_concat,
            X_test=X_test,
            y_test=y_test,
            pos_label=self._pos_label,
        )
        return report

    @property
    def pos_label(self) -> PositiveLabel | None:
        return self._pos_label

    ####################################################################################
    # Methods related to the help and repr
    ####################################################################################

    def _get_help_title(self) -> str:
        return "Tools to compare estimators"

    def _get_help_legend(self) -> str:
        return (
            "[cyan](↗︎)[/cyan] higher is better [orange1](↘︎)[/orange1] lower is better"
        )

    def __repr__(self) -> str:
        """Return a string representation."""
        return f"{self.__class__.__name__}(...)"

    def _repr_html_(self) -> str:
        """HTML representation with a selector to inspect one compared report."""
        metrics_html = self.metrics.summarize(data_source="test").frame()._repr_html_()

        comparison_reports = []
        for label, report in self.reports_.items():
            fragments = report._html_repr_fragments()
            comparison_reports.append(
                {
                    "label": label,
                    "estimator_display": fragments["estimator_display"],
                    "table_report": fragments["table_report"],
                }
            )

        container_id = f"skore-comparison-report-{uuid.uuid4().hex[:8]}"
        help_doc_url = get_documentation_url(obj=self, method_name="help")
        report_class_name = self.__class__.__name__
        metrics_accessor_doc_url = get_documentation_url(
            obj=self, accessor_name="metrics"
        )
        inspection_accessor_doc_url = get_documentation_url(
            obj=self, accessor_name="inspection"
        )
        return render_template(
            "comparison_report.html.j2",
            {
                "container_id": container_id,
                "metrics_summary": metrics_html,
                "comparison_reports": comparison_reports,
                "help_doc_url": help_doc_url,
                "report_class_name": report_class_name,
                "metrics_accessor_doc_url": metrics_accessor_doc_url,
                "inspection_accessor_doc_url": inspection_accessor_doc_url,
            },
        )

    def _repr_mimebundle_(self, **kwargs):
        """Mime bundle used by Jupyter kernels to display the report."""
        output = {"text/plain": repr(self)}
        output["text/html"] = self._repr_html_()
        return output


def _deduplicate_report_names(report_names: list[str]) -> list[str]:
    """De-duplicate report names that appear several times.

    Leave the other report names alone.

    Parameters
    ----------
    report_names : list of str
        The list of report names to be checked.

    Returns
    -------
    list of str
        The de-duplicated list of report names.

    Examples
    --------
    >>> _deduplicate_report_names(['a', 'b'])
    ['a', 'b']
    >>> _deduplicate_report_names(['a', 'a'])
    ['a_1', 'a_2']
    >>> _deduplicate_report_names(['a', 'b', 'a'])
    ['a_1', 'b', 'a_2']
    >>> _deduplicate_report_names(['a', 'b', 'a', 'b'])
    ['a_1', 'b_1', 'a_2', 'b_2']
    >>> _deduplicate_report_names([])
    []
    >>> _deduplicate_report_names(['a'])
    ['a']
    """
    counts = Counter(report_names)
    if len(report_names) == len(counts):
        return report_names

    names = report_names.copy()
    seen: Counter = Counter()
    for i in range(len(names)):
        name = names[i]
        seen[name] += 1
        if counts[name] > 1:
            names[i] = f"{name}_{seen[name]}"
    return names
