#!/usr/bin/env python3 -u
# -*- coding: utf-8 -*-
# copyright: sktime developers, BSD-3-Clause License (see LICENSE file)
"""Implements grid search functionality to tune forecasters."""

__author__ = ["mloning"]
__all__ = ["ForecastingGridSearchCV", "ForecastingRandomizedSearchCV"]

from collections.abc import Sequence

import numpy as np
import pandas as pd
from joblib import Parallel, delayed
from sklearn.model_selection import ParameterGrid, ParameterSampler, check_cv
from sklearn.utils.metaestimators import if_delegate_has_method

from sktime.exceptions import NotFittedError
from sktime.forecasting.base import BaseForecaster
from sktime.forecasting.model_evaluation import evaluate
from sktime.utils.validation.forecasting import check_scoring


class BaseGridSearch(BaseForecaster):

    _tags = {
        "scitype:y": "both",
        "requires-fh-in-fit": False,
        "handles-missing-data": False,
        "ignores-exogeneous-X": True,
        "capability:pred_int": True,
    }

    def __init__(
        self,
        forecaster,
        cv,
        strategy="refit",
        n_jobs=None,
        pre_dispatch=None,
        backend="loky",
        refit=False,
        scoring=None,
        verbose=0,
        return_n_best_forecasters=1,
    ):

        self.forecaster = forecaster
        self.cv = cv
        self.strategy = strategy
        self.n_jobs = n_jobs
        self.pre_dispatch = pre_dispatch
        self.backend = backend
        self.refit = refit
        self.scoring = scoring
        self.verbose = verbose
        self.return_n_best_forecasters = return_n_best_forecasters
        super(BaseGridSearch, self).__init__()
        tags_to_clone = [
            "requires-fh-in-fit",
            "capability:pred_int",
            "scitype:y",
            "ignores-exogeneous-X",
            "handles-missing-data",
            "y_inner_mtype",
            "X_inner_mtype",
            "X-y-must-have-same-index",
            "enforce_index_type",
        ]
        self.clone_tags(forecaster, tags_to_clone)

    @if_delegate_has_method(delegate=("best_forecaster_", "forecaster"))
    def _update(self, y, X=None, update_params=False):
        """Call _update on the forecaster with the best found parameters."""
        self.check_is_fitted("update")
        self.best_forecaster_._update(y, X, update_params=update_params)
        return self

    @if_delegate_has_method(delegate=("best_forecaster_", "forecaster"))
    def _predict(self, fh=None, X=None):
        """Call _predict on the forecaster with the best found parameters."""
        self.check_is_fitted("predict")
        return self.best_forecaster_._predict(fh=fh, X=X)

    @if_delegate_has_method(delegate=("best_forecaster_", "forecaster"))
    def _predict_interval(self, fh=None, X=None, coverage=None):
        """Call _predict_interval on the forecaster with the best found parameters."""
        self.check_is_fitted("predict")
        return self.best_forecaster_._predict_interval(fh=fh, X=X, coverage=coverage)

    @if_delegate_has_method(delegate=("best_forecaster_", "forecaster"))
    def _predict_quantiles(self, fh=None, X=None, alpha=None):
        """Call _predict_quantiles on the forecaster with the best found parameters."""
        self.check_is_fitted("predict")
        return self.best_forecaster_._predict_quantiles(fh=fh, X=X, alpha=alpha)

    @if_delegate_has_method(delegate=("best_forecaster_", "forecaster"))
    def _predict_var(self, fh=None, X=None, cov=False):
        """Call _predict_var on the forecaster with the best found parameters."""
        self.check_is_fitted("predict")
        return self.best_forecaster_._predict_var(fh=fh, X=X, cov=cov)

    @if_delegate_has_method(delegate=("best_forecaster_", "forecaster"))
    def _predict_proba(self, fh=None, X=None, marginal=True):
        """Call _predict_proba on the forecaster with the best found parameters."""
        self.check_is_fitted("predict")
        return self.best_forecaster_._predict_proba(fh=fh, X=X, marginal=marginal)

    @if_delegate_has_method(delegate=("best_forecaster_", "forecaster"))
    def transform(self, y, X=None):
        """Call transform on the forecaster with the best found parameters."""
        self.check_is_fitted("transform")
        return self.best_forecaster_.transform(y, X)

    @if_delegate_has_method(delegate=("best_forecaster_", "forecaster"))
    def get_fitted_params(self):
        """Get fitted parameters.

        Returns
        -------
        fitted_params : dict
        """
        self.check_is_fitted("get_fitted_params")
        return self.best_forecaster_.get_fitted_params()

    @if_delegate_has_method(delegate=("best_forecaster_", "forecaster"))
    def inverse_transform(self, y, X=None):
        """Call inverse_transform on the forecaster with the best found params.

        Only available if the underlying forecaster implements
        ``inverse_transform`` and ``refit=True``.

        Parameters
        ----------
        y : indexable, length n_samples
            Must fulfill the input assumptions of the
            underlying forecaster.
        """
        self.check_is_fitted("inverse_transform")
        return self.best_forecaster_.inverse_transform(y, X)

    def score(self, y, X=None, fh=None):
        """Return the score on the given data, if forecaster been refitted.

        This uses the score defined by ``scoring`` where provided, and the
        ``best_forecaster_.score`` method otherwise.

        Parameters
        ----------
        y : pandas.Series
            Target time series to which to compare the forecasts.
        X : pandas.DataFrame, shape=[n_obs, n_vars], optional (default=None)
            An optional 2-d dataframe of exogenous variables.
        fh : ForecastingHorizon, int, np.ndarray, pd.Index, optional (default=None)
            Forecasting horizon

        Returns
        -------
        score : float
        """
        self.check_is_fitted("score")

        if self.scoring is None:
            return self.best_forecaster_.score(y, X=X, fh=fh)

        else:
            y_pred = self.best_forecaster_.predict(fh, X=X)
            return self.scoring(y, y_pred)

    def _run_search(self, evaluate_candidates):
        raise NotImplementedError("abstract method")

    def check_is_fitted(self, method_name=None):
        """Check if `fit` has been called.

        Parameters
        ----------
        method_name : str
            Name of the calling method.

        Raises
        ------
        NotFittedError
            If forecaster has not been fitted yet.
        """
        super(BaseGridSearch, self).check_is_fitted()

        # We additionally check if the tuned forecaster has been fitted.
        if method_name is not None:
            if not self.refit:
                raise NotFittedError(
                    "This %s instance was initialized "
                    "with refit=False. %s is "
                    "available only after refitting on the "
                    "best parameters. You can refit an forecaster "
                    "manually using the ``best_params_`` "
                    "attribute" % (type(self).__name__, method_name)
                )
            else:
                self.best_forecaster_.check_is_fitted()

    def _fit(self, y, X=None, fh=None):
        """Fit to training data.

        Parameters
        ----------
        y : pd.Series
            Target time series to which to fit the forecaster.
        fh : int, list or np.array, optional (default=None)
            The forecasters horizon with the steps ahead to to predict.
        X : pd.DataFrame, optional (default=None)
            Exogenous variables are ignored

        Returns
        -------
        self : returns an instance of self.
        """
        cv = check_cv(self.cv)

        scoring = check_scoring(self.scoring)
        scoring_name = f"test_{scoring.name}"

        parallel = Parallel(
            n_jobs=self.n_jobs, pre_dispatch=self.pre_dispatch, backend=self.backend
        )

        def _fit_and_score(params):
            # Clone forecaster.
            forecaster = self.forecaster.clone()

            # Set parameters.
            forecaster.set_params(**params)

            # Evaluate.
            out = evaluate(
                forecaster,
                cv,
                y,
                X,
                strategy=self.strategy,
                scoring=scoring,
            )

            # Filter columns.
            out = out.filter(items=[scoring_name, "fit_time", "pred_time"], axis=1)

            # Aggregate results.
            out = out.mean()
            out = out.add_prefix("mean_")

            # Add parameters to output table.
            out["params"] = params

            return out

        def evaluate_candidates(candidate_params):
            candidate_params = list(candidate_params)

            if self.verbose > 0:
                n_candidates = len(candidate_params)
                n_splits = cv.get_n_splits(y)
                print(  # noqa
                    "Fitting {0} folds for each of {1} candidates,"
                    " totalling {2} fits".format(
                        n_splits, n_candidates, n_candidates * n_splits
                    )
                )

            out = parallel(
                delayed(_fit_and_score)(params) for params in candidate_params
            )

            if len(out) < 1:
                raise ValueError(
                    "No fits were performed. "
                    "Was the CV iterator empty? "
                    "Were there no candidates?"
                )

            return out

        # Run grid-search cross-validation.
        results = self._run_search(evaluate_candidates)

        results = pd.DataFrame(results)

        # Rank results, according to whether greater is better for the given scoring.
        results[f"rank_{scoring_name}"] = results.loc[:, f"mean_{scoring_name}"].rank(
            ascending=not scoring.greater_is_better
        )

        self.cv_results_ = results

        # Select best parameters.
        self.best_index_ = results.loc[:, f"rank_{scoring_name}"].argmin()
        self.best_score_ = results.loc[self.best_index_, f"mean_{scoring_name}"]
        self.best_params_ = results.loc[self.best_index_, "params"]
        self.best_forecaster_ = self.forecaster.clone().set_params(**self.best_params_)

        # Refit model with best parameters.
        if self.refit:
            self.best_forecaster_.fit(y, X, fh)

        # Sort values according to rank
        results = results.sort_values(
            by=f"rank_{scoring_name}", ascending=not scoring.greater_is_better
        )
        # Select n best forecaster
        self.n_best_forecasters_ = []
        self.n_best_scores_ = []
        for i in range(self.return_n_best_forecasters):
            params = results["params"].iloc[i]
            rank = results[f"rank_{scoring_name}"].iloc[i]
            rank = str(int(rank))
            forecaster = self.forecaster.clone().set_params(**params)
            # Refit model with best parameters.
            if self.refit:
                forecaster.fit(y, X, fh)
            self.n_best_forecasters_.append((rank, forecaster))
            # Save score
            score = results[f"mean_{scoring_name}"].iloc[i]
            self.n_best_scores_.append(score)

        return self


class ForecastingGridSearchCV(BaseGridSearch):
    """Perform grid-search cross-validation to find optimal model parameters.

    The forecaster is fit on the initial window and then temporal
    cross-validation is used to find the optimal parameter

    Grid-search cross-validation is performed based on a cross-validation
    iterator encoding the cross-validation scheme, the parameter grid to
    search over, and (optionally) the evaluation metric for comparing model
    performance. As in scikit-learn, tuning works through the common
    hyper-parameter interface which allows to repeatedly fit and evaluate
    the same forecaster with different hyper-parameters.

    Parameters
    ----------
    forecaster : estimator object
        The estimator should implement the sktime or scikit-learn estimator
        interface. Either the estimator must contain a "score" function,
        or a scoring function must be passed.
    cv : cross-validation generator or an iterable
        e.g. SlidingWindowSplitter()
    param_grid : dict or list of dictionaries
        Model tuning parameters of the forecaster to evaluate
    scoring: function, optional (default=None)
        Function to score models for evaluation of optimal parameters
    n_jobs: int, optional (default=None)
        Number of jobs to run in parallel.
        None means 1 unless in a joblib.parallel_backend context.
        -1 means using all processors.
    refit: bool, optional (default=True)
        Refit the forecaster with the best parameters on all the data
    verbose: int, optional (default=0)
    return_n_best_forecasters: int, default=1
        In case the n best forecaster should be returned, this value can be set
        and the n best forecasters will be assigned to n_best_forecasters_
    pre_dispatch: str, optional (default='2*n_jobs')
    error_score: numeric value or the str 'raise', optional (default=np.nan)
        The test score returned when a forecaster fails to be fitted.
    return_train_score: bool, optional (default=False)
    backend: str, optional (default="loky")
        Specify the parallelisation backend implementation in joblib, where
        "loky" is used by default.

    Attributes
    ----------
    best_index_ : int
    best_score_: float
        Score of the best model
    best_params_ : dict
        Best parameter values across the parameter grid
    best_forecaster_ : estimator
        Fitted estimator with the best parameters
    cv_results_ : dict
        Results from grid search cross validation
    n_splits_: int
        Number of splits in the data for cross validation}
    refit_time_ : float
        Time (seconds) to refit the best forecaster
    scorer_ : function
        Function used to score model
    n_best_forecasters_: list of tuples ("rank", <forecaster>)
        The "rank" is in relation to best_forecaster_
    n_best_scores_: list of float
        The scores of n_best_forecasters_ sorted from best to worst
        score of forecasters

    Examples
    --------
    >>> from sktime.datasets import load_shampoo_sales
    >>> from sktime.forecasting.model_selection import (
    ...     ExpandingWindowSplitter,
    ...     ForecastingGridSearchCV,
    ...     ExpandingWindowSplitter)
    >>> from sktime.forecasting.naive import NaiveForecaster
    >>> y = load_shampoo_sales()
    >>> fh = [1,2,3]
    >>> cv = ExpandingWindowSplitter(
    ...     start_with_window=True,
    ...     fh=fh)
    >>> forecaster = NaiveForecaster()
    >>> param_grid = {"strategy" : ["last", "mean", "drift"]}
    >>> gscv = ForecastingGridSearchCV(
    ...     forecaster=forecaster,
    ...     param_grid=param_grid,
    ...     cv=cv)
    >>> gscv.fit(y)
    ForecastingGridSearchCV(...)
    >>> y_pred = gscv.predict(fh)

        Advanced model meta-tuning (model selection) with multiple forecasters
        together with hyper-parametertuning at same time using sklearn notation:
    >>> from sktime.datasets import load_shampoo_sales
    >>> from sktime.forecasting.exp_smoothing import ExponentialSmoothing
    >>> from sktime.forecasting.naive import NaiveForecaster
    >>> from sktime.forecasting.model_selection import ExpandingWindowSplitter
    >>> from sktime.forecasting.model_selection import ForecastingGridSearchCV
    >>> from sktime.forecasting.compose import TransformedTargetForecaster
    >>> from sktime.forecasting.theta import ThetaForecaster
    >>> from sktime.transformations.series.impute import Imputer
    >>> y = load_shampoo_sales()
    >>> pipe = TransformedTargetForecaster(steps=[
    ...     ("imputer", Imputer()),
    ...     ("forecaster", NaiveForecaster())])
    >>> cv = ExpandingWindowSplitter(
    ...     initial_window=24,
    ...     step_length=12,
    ...     start_with_window=True,
    ...     fh=[1,2,3])
    >>> gscv = ForecastingGridSearchCV(
    ...     forecaster=pipe,
    ...     param_grid=[{
    ...         "forecaster": [NaiveForecaster(sp=12)],
    ...         "forecaster__strategy": ["drift", "last", "mean"],
    ...     },
    ...     {
    ...         "imputer__method": ["mean", "drift"],
    ...         "forecaster": [ThetaForecaster(sp=12)],
    ...     },
    ...     {
    ...         "imputer__method": ["mean", "median"],
    ...         "forecaster": [ExponentialSmoothing(sp=12)],
    ...         "forecaster__trend": ["add", "mul"],
    ...     },
    ...     ],
    ...     cv=cv,
    ...     n_jobs=-1)
    >>> gscv.fit(y)
    ForecastingGridSearchCV(...)
    >>> y_pred = gscv.predict(fh=[1,2,3])
    """

    _required_parameters = ["forecaster", "cv", "param_grid"]

    def __init__(
        self,
        forecaster,
        cv,
        param_grid,
        scoring=None,
        strategy="refit",
        n_jobs=None,
        refit=True,
        verbose=0,
        return_n_best_forecasters=1,
        pre_dispatch="2*n_jobs",
        backend="loky",
    ):
        super(ForecastingGridSearchCV, self).__init__(
            forecaster=forecaster,
            scoring=scoring,
            n_jobs=n_jobs,
            refit=refit,
            cv=cv,
            strategy=strategy,
            verbose=verbose,
            return_n_best_forecasters=return_n_best_forecasters,
            pre_dispatch=pre_dispatch,
            backend=backend,
        )
        self.param_grid = param_grid

    def _check_param_grid(self, param_grid):
        """_check_param_grid from sklearn 1.0.2, before it was removed."""
        if hasattr(param_grid, "items"):
            param_grid = [param_grid]

        for p in param_grid:
            for name, v in p.items():
                if isinstance(v, np.ndarray) and v.ndim > 1:
                    raise ValueError("Parameter array should be one-dimensional.")

                if isinstance(v, str) or not isinstance(v, (np.ndarray, Sequence)):
                    raise ValueError(
                        "Parameter grid for parameter ({0}) needs to"
                        " be a list or numpy array, but got ({1})."
                        " Single values need to be wrapped in a list"
                        " with one element.".format(name, type(v))
                    )

                if len(v) == 0:
                    raise ValueError(
                        "Parameter values for parameter ({0}) need "
                        "to be a non-empty sequence.".format(name)
                    )

    def _run_search(self, evaluate_candidates):
        """Search all candidates in param_grid."""
        self._check_param_grid(self.param_grid)
        return evaluate_candidates(ParameterGrid(self.param_grid))

    @classmethod
    def get_test_params(cls, parameter_set="default"):
        """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.

        Returns
        -------
        params : dict or list of dict
        """
        from sktime.forecasting.model_selection._split import SingleWindowSplitter
        from sktime.forecasting.naive import NaiveForecaster
        from sktime.performance_metrics.forecasting import MeanAbsolutePercentageError

        params = {
            "forecaster": NaiveForecaster(strategy="mean"),
            "cv": SingleWindowSplitter(fh=1),
            "param_grid": {"window_length": [2, 5]},
            "scoring": MeanAbsolutePercentageError(symmetric=True),
        }
        return params


class ForecastingRandomizedSearchCV(BaseGridSearch):
    """Perform randomized-search cross-validation to find optimal model parameters.

    The forecaster is fit on the initial window and then temporal
    cross-validation is used to find the optimal parameter

    Randomized cross-validation is performed based on a cross-validation
    iterator encoding the cross-validation scheme, the parameter distributions to
    search over, and (optionally) the evaluation metric for comparing model
    performance. As in scikit-learn, tuning works through the common
    hyper-parameter interface which allows to repeatedly fit and evaluate
    the same forecaster with different hyper-parameters.

    Parameters
    ----------
    forecaster : estimator object
        The estimator should implement the sktime or scikit-learn estimator
        interface. Either the estimator must contain a "score" function,
        or a scoring function must be passed.
    cv : cross-validation generator or an iterable
        e.g. SlidingWindowSplitter()
    param_distributions : dict or list of dicts
        Dictionary with parameters names (`str`) as keys and distributions
        or lists of parameters to try. Distributions must provide a ``rvs``
        method for sampling (such as those from scipy.stats.distributions).
        If a list is given, it is sampled uniformly.
        If a list of dicts is given, first a dict is sampled uniformly, and
        then a parameter is sampled using that dict as above.
    n_iter : int, default=10
        Number of parameter settings that are sampled. n_iter trades
        off runtime vs quality of the solution.
    scoring: function, optional (default=None)
        Function to score models for evaluation of optimal parameters
    n_jobs: int, optional (default=None)
        Number of jobs to run in parallel.
        None means 1 unless in a joblib.parallel_backend context.
        -1 means using all processors.
    refit: bool, optional (default=True)
        Refit the forecaster with the best parameters on all the data
    verbose: int, optional (default=0)
    return_n_best_forecasters: int, default=1
        In case the n best forecaster should be returned, this value can be set
        and the n best forecasters will be assigned to n_best_forecasters_
    pre_dispatch: str, optional (default='2*n_jobs')
    random_state : int, RandomState instance or None, default=None
        Pseudo random number generator state used for random uniform sampling
        from lists of possible values instead of scipy.stats distributions.
        Pass an int for reproducible output across multiple
        function calls.
    pre_dispatch: str, optional (default='2*n_jobs')
    backend: str, optional (default="loky")
        Specify the parallelisation backend implementation in joblib, where
        "loky" is used by default.

    Attributes
    ----------
    best_index_ : int
    best_score_: float
        Score of the best model
    best_params_ : dict
        Best parameter values across the parameter grid
    best_forecaster_ : estimator
        Fitted estimator with the best parameters
    cv_results_ : dict
        Results from grid search cross validation
    n_best_forecasters_: list of tuples ("rank", <forecaster>)
        The "rank" is in relation to best_forecaster_
    n_best_scores_: list of float
        The scores of n_best_forecasters_ sorted from best to worst
        score of forecasters
    """

    _required_parameters = ["forecaster", "cv", "param_distributions"]

    def __init__(
        self,
        forecaster,
        cv,
        param_distributions,
        n_iter=10,
        scoring=None,
        strategy="refit",
        n_jobs=None,
        refit=True,
        verbose=0,
        return_n_best_forecasters=1,
        random_state=None,
        pre_dispatch="2*n_jobs",
        backend="loky",
    ):
        super(ForecastingRandomizedSearchCV, self).__init__(
            forecaster=forecaster,
            scoring=scoring,
            strategy=strategy,
            n_jobs=n_jobs,
            refit=refit,
            cv=cv,
            verbose=verbose,
            return_n_best_forecasters=return_n_best_forecasters,
            pre_dispatch=pre_dispatch,
            backend=backend,
        )
        self.param_distributions = param_distributions
        self.n_iter = n_iter
        self.random_state = random_state

    def _run_search(self, evaluate_candidates):
        """Search n_iter candidates from param_distributions."""
        return evaluate_candidates(
            ParameterSampler(
                self.param_distributions, self.n_iter, random_state=self.random_state
            )
        )

    @classmethod
    def get_test_params(cls, parameter_set="default"):
        """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.

        Returns
        -------
        params : dict or list of dict
        """
        from sktime.forecasting.model_selection._split import SingleWindowSplitter
        from sktime.forecasting.naive import NaiveForecaster
        from sktime.performance_metrics.forecasting import MeanAbsolutePercentageError

        params = {
            "forecaster": NaiveForecaster(strategy="mean"),
            "cv": SingleWindowSplitter(fh=1),
            "param_distributions": {"window_length": [2, 5]},
            "scoring": MeanAbsolutePercentageError(symmetric=True),
        }
        return params
