"""A shapelet transform classifier (STC).

Shapelet transform classifier pipeline that simply performs a (configurable) shapelet
transform then builds (by default) a rotation forest classifier on the output.
"""

__maintainer__ = ["TonyBagnall"]
__all__ = ["ShapeletTransformClassifier"]


import numpy as np
from sklearn.model_selection import cross_val_predict
from sklearn.utils import check_random_state

from aeon.base._base import _clone_estimator
from aeon.classification.base import BaseClassifier
from aeon.classification.sklearn import RotationForestClassifier
from aeon.transformations.collection.shapelet_based import RandomShapeletTransform
from aeon.utils.validation import check_n_jobs


class ShapeletTransformClassifier(BaseClassifier):
    """
    A shapelet transform classifier (STC).

    Implementation of the binary shapelet transform classifier pipeline along the lines
    of [1]_, [2]_, but with random shapelet sampling. Transforms the data using the
    configurable `RandomShapeletTransform` and then builds a `RotationForestClassifier`
    classifier.

    As some implementations and applications contract the transformation solely,
    contracting is available for the transform only and both classifier and transform.

    Parameters
    ----------
    n_shapelet_samples : int, default=10000
        The number of candidate shapelets to be considered for the final transform.
        Filtered down to ``<= max_shapelets``, keeping the shapelets with the most
        information gain.
    max_shapelets : int or None, default=None
        Max number of shapelets to keep for the final transform. Each class value will
        have its own max, set to ``n_classes_ / max_shapelets``. If `None`, uses the
        minimum between ``10 * n_cases_`` and `1000`.
    max_shapelet_length : int or None, default=None
        Lower bound on candidate shapelet lengths for the transform. If ``None``, no
        max length is used
    estimator : BaseEstimator or None, default=None
        Base estimator for the ensemble, can be supplied a sklearn `BaseEstimator`. If
        `None` a default `RotationForestClassifier` classifier is used.
    batch_size : int or None, default=100
        Number of shapelet candidates processed before being merged into the set of best
        shapelets in the transform.
    verbose : bool, default=False
        Whether to print progress messages during fitting and transforming.
    transform_limit_in_minutes : int, default=0
        Time contract to limit transform time in minutes for the shapelet transform,
        overriding `n_shapelet_samples`. A value of `0` means ``n_shapelet_samples``
        is used.
    time_limit_in_minutes : int, default=0
        Time contract to limit build time in minutes, overriding ``n_shapelet_samples``
        and ``transform_limit_in_minutes``. The ``estimator`` will only be contracted if
        a ``time_limit_in_minutes parameter`` is present. Default of `0` means
        ``n_shapelet_samples`` or ``transform_limit_in_minutes`` is used.
    contract_max_n_shapelet_samples : int, default=np.inf
        Max number of shapelets to extract when contracting the transform with
        ``transform_limit_in_minutes`` or ``time_limit_in_minutes``.
    n_jobs : int, default=1
        The number of jobs to run in parallel for both ``fit`` and ``predict``.
        `-1` means using all processors.
    random_state : int, RandomState instance or None, default=None
        If `int`, random_state is the seed used by the random number generator;
        If `RandomState` instance, random_state is the random number generator;
        If `None`, the random number generator is the `RandomState` instance used
        by `np.random`.

    Attributes
    ----------
    classes_ : list
        The unique class labels in the training set.
    n_classes_ : int
        The number of unique classes in the training set.
    n_cases_ : int
        The number of train cases in the training set.
    n_channels_ : int
        The number of dimensions per case in the training set.

    See Also
    --------
    RandomShapeletTransform : The randomly sampled shapelet transform.
    RotationForestClassifier : The default rotation forest classifier used.

    Notes
    -----
    For the Java version, see
    `tsml <https://github.com/uea-machine-learning/tsml/blob/master/src/main/
    java/tsml/classifiers/shapelet_based/ShapeletTransformClassifier.java>`_.

    References
    ----------
    .. [1] Jon Hills et al., "Classification of time series by shapelet transformation",
       Data Mining and Knowledge Discovery, 28(4), 851--881, 2014.
    .. [2] A. Bostrom and A. Bagnall, "Binary Shapelet Transform for Multiclass Time
       Series Classification", Transactions on Large-Scale Data and Knowledge Centered
       Systems, 32, 2017.

    Examples
    --------
    >>> from aeon.classification.shapelet_based import ShapeletTransformClassifier
    >>> from aeon.classification.sklearn import RotationForestClassifier
    >>> from aeon.datasets import load_unit_test
    >>> X_train, y_train = load_unit_test(split="train")
    >>> X_test, y_test = load_unit_test(split="test")
    >>> clf = ShapeletTransformClassifier(
    ...     estimator=RotationForestClassifier(n_estimators=3),
    ...     n_shapelet_samples=100,
    ...     max_shapelets=10,
    ...     batch_size=20,
    ... )
    >>> clf.fit(X_train, y_train)
    ShapeletTransformClassifier(...)
    >>> y_pred = clf.predict(X_test)
    """

    _tags = {
        "capability:multivariate": True,
        "capability:train_estimate": True,
        "capability:contractable": True,
        "capability:multithreading": True,
        "capability:unequal_length": True,
        "algorithm_type": "shapelet",
        "X_inner_type": ["np-list", "numpy3D"],
    }

    def __init__(
        self,
        n_shapelet_samples: int = 10000,
        max_shapelets: int | None = None,
        max_shapelet_length: int | None = None,
        estimator=None,
        batch_size: int | None = 100,
        verbose: bool = False,
        transform_limit_in_minutes: int = 0,
        time_limit_in_minutes: int = 0,
        contract_max_n_shapelet_samples: int = np.inf,
        n_jobs: int = 1,
        random_state: int | np.random.RandomState | None = None,
    ) -> None:
        self.n_shapelet_samples = n_shapelet_samples
        self.max_shapelets = max_shapelets
        self.max_shapelet_length = max_shapelet_length
        self.estimator = estimator
        self.batch_size = batch_size
        self.verbose = verbose
        self.transform_limit_in_minutes = transform_limit_in_minutes
        self.time_limit_in_minutes = time_limit_in_minutes
        self.contract_max_n_shapelet_samples = contract_max_n_shapelet_samples
        self.random_state = random_state
        self.n_jobs = n_jobs

        super().__init__()

    def _fit(self, X, y):
        """Fit ShapeletTransformClassifier to training data.

        Parameters
        ----------
        X : 3D np.ndarray of shape = [n_cases, n_channels, n_timepoints]
            The training data.
        y : array-like, shape = [n_cases]
            The class labels.

        Returns
        -------
        self :
            Reference to self.

        Notes
        -----
        Changes state by creating a fitted model that updates attributes
        ending in "_".
        """
        X_t = self._fit_stc_shared(X, y)

        if self.verbose:
            print("Fitting estimator...")  # noqa: T201

        self._estimator.fit(X_t, y)

        if self.verbose:
            print("Finished fitting estimator...")  # noqa: T201

    def _predict(self, X) -> np.ndarray:
        """Predicts labels for sequences in X.

        Parameters
        ----------
        X : 3D np.ndarray of shape = [n_cases, n_channels, n_timepoints]
            The data to make predictions for.

        Returns
        -------
        y : array-like, shape = [n_cases]
            Predicted class labels.
        """
        if self.verbose:
            print("Transforming predict X...")  # noqa: T201

        X_t = self._transformer.transform(X)
        X_t = np.nan_to_num(X_t, False, -1, -1, -1)

        if self.verbose:
            print("Finished transforming predict X...")  # noqa: T201
            print("Predicting...")  # noqa: T201

        pred = self._estimator.predict(X_t)

        if self.verbose:
            print("Finished predicting...")  # noqa: T201

        return pred

    def _predict_proba(self, X) -> np.ndarray:
        """Predicts labels probabilities for sequences in X.

        Parameters
        ----------
        X : 3D np.ndarray of shape = [n_cases, n_channels, n_timepoints]
            The data to make predict probabilities for.

        Returns
        -------
        y : array-like, shape = [n_cases, n_classes_]
            Predicted probabilities using the ordering in classes_.
        """
        if self.verbose:
            print("Transforming predict_proba X...")  # noqa: T201

        X_t = self._transformer.transform(X)
        X_t = np.nan_to_num(X_t, False, -1, -1, -1)

        if self.verbose:
            print("Finished transforming predict_proba X...")  # noqa: T201
            print("Predicting probabilities...")  # noqa: T201

        m = getattr(self._estimator, "predict_proba", None)
        if callable(m):
            proba = self._estimator.predict_proba(X_t)
        else:
            proba = np.zeros((X.shape[0], self.n_classes_))
            preds = self._estimator.predict(X_t)
            for i in range(0, X.shape[0]):
                proba[i, np.where(self.classes_ == preds[i])] = 1

        if self.verbose:
            print("Finished predicting probabilities...")  # noqa: T201

        return proba

    def _fit_predict(self, X, y) -> np.ndarray:
        rng = check_random_state(self.random_state)
        return np.array(
            [
                self.classes_[int(rng.choice(np.flatnonzero(prob == prob.max())))]
                for prob in self._fit_predict_proba(X, y)
            ]
        )

    def _fit_predict_proba(self, X, y) -> np.ndarray:
        X_t = self._fit_stc_shared(X, y)

        if (isinstance(self.estimator, RotationForestClassifier)) or (
            self.estimator is None
        ):
            if self.verbose:
                print(  # noqa: T201
                    "Fitting estimator and generating train set estimates "
                    "(RotF OOB)..."
                )

            proba = self._estimator.fit_predict_proba(X_t, y)
        else:
            if self.verbose:
                print(  # noqa: T201
                    "Fitting estimator and generating train set estimates "
                    "(10 fold CV)..."
                )

            self._estimator.fit(X_t, y)

            m = getattr(self._estimator, "predict_proba", None)
            if not callable(m):
                raise ValueError("Estimator must have a predict_proba method.")

            cv_size = 10
            _, counts = np.unique(y, return_counts=True)
            min_class = np.min(counts)
            if min_class < cv_size:
                cv_size = min_class
                if cv_size < 2:
                    raise ValueError(
                        "All classes must have at least 2 values to run the "
                        "fit_predict/fit_predict_proba cross-validation."
                    )

            estimator = _clone_estimator(self.estimator, self.random_state)

            proba = cross_val_predict(
                estimator,
                X=X_t,
                y=y,
                cv=cv_size,
                method="predict_proba",
                n_jobs=self._n_jobs,
            )

        if self.verbose:
            print(  # noqa: T201
                "Finished fitting estimator and generating train set estimates..."
            )

        return proba

    def _fit_stc_shared(self, X, y):
        self.n_instances_ = len(X)
        self.n_channels_ = X[0].shape[0]
        self._n_jobs = check_n_jobs(self.n_jobs)

        self._transform_limit_in_minutes = 0
        self._classifier_limit_in_minutes = 0
        if self.time_limit_in_minutes > 0:
            # contracting 2/3 transform (with 1/5 of that taken away for final
            # transform), 1/3 classifier
            third = self.time_limit_in_minutes / 3
            self._classifier_limit_in_minutes = third
            self._transform_limit_in_minutes = (third * 2) / 5 * 4
        elif self.transform_limit_in_minutes > 0:
            self._transform_limit_in_minutes = self.transform_limit_in_minutes

        self._transformer = RandomShapeletTransform(
            n_shapelet_samples=self.n_shapelet_samples,
            max_shapelets=self.max_shapelets,
            max_shapelet_length=self.max_shapelet_length,
            batch_size=self.batch_size,
            verbose=self.verbose,
            time_limit_in_minutes=self._transform_limit_in_minutes,
            contract_max_n_shapelet_samples=self.contract_max_n_shapelet_samples,
            n_jobs=self.n_jobs,
            random_state=self.random_state,
        )

        self._estimator = _clone_estimator(
            RotationForestClassifier() if self.estimator is None else self.estimator,
            self.random_state,
        )

        m = getattr(self._estimator, "n_jobs", None)
        if m is not None:
            self._estimator.n_jobs = self._n_jobs

        m = getattr(self._estimator, "time_limit_in_minutes", None)
        if m is not None and self.time_limit_in_minutes > 0:
            self._estimator.time_limit_in_minutes = self._classifier_limit_in_minutes

        if self.verbose:
            print("Fitting and transforming shapelets...")  # noqa: T201

        X_t = self._transformer.fit_transform(X, y)
        X_t = np.nan_to_num(X_t, False, -1, -1, -1)

        if self.verbose:
            print("Finished and transforming shapelets...")  # noqa: T201

        return X_t

    @classmethod
    def _get_test_params(cls, parameter_set: str = "default") -> dict | list[dict]:
        """Return testing parameter settings for the estimator.

        Parameters
        ----------
        parameter_set : str, default="default"
            Name of the set of test parameters to return, for use in tests. If no
            special parameters are defined for a value, will return `"default"` set.
            ShapeletTransformClassifier provides the following special sets:
                 "results_comparison" - used in some classifiers to compare against
                    previously generated results where the default set of parameters
                    cannot produce suitable probability estimates
                "contracting" - used in classifiers that set the
                    "capability:contractable" tag to True to test contacting
                    functionality
                "train_estimate" - used in some classifiers that set the
                    "capability:train_estimate" tag to True to allow for more efficient
                    testing when relevant parameters are available

        Returns
        -------
        params : dict or list of dict, default={}
            Parameters to create testing instances of the class.
            Each dict are parameters to construct an "interesting" test instance, i.e.,
            `MyClass(**params)` or `MyClass(**params[i])` creates a valid test instance.
        """
        from sklearn.ensemble import RandomForestClassifier

        if parameter_set == "results_comparison":
            return {
                "estimator": RandomForestClassifier(n_estimators=5),
                "n_shapelet_samples": 50,
                "max_shapelets": 10,
                "batch_size": 10,
            }
        elif parameter_set == "contracting":
            return {
                "time_limit_in_minutes": 5,
                "estimator": RotationForestClassifier(contract_max_n_estimators=2),
                "contract_max_n_shapelet_samples": 10,
                "max_shapelets": 3,
                "batch_size": 5,
            }
        else:
            return {
                "estimator": RotationForestClassifier(n_estimators=2),
                "n_shapelet_samples": 10,
                "max_shapelets": 3,
                "batch_size": 5,
            }
