#!/usr/bin/env python3 -u
# -*- coding: utf-8 -*-
# copyright: sktime developers, BSD-3-Clause License (see LICENSE file)

__author__ = ["Markus Löning"]
__all__ = ["ForecastingGridSearchCV", "ForecastingRandomizedSearchCV"]

import numbers
import time
import warnings
from collections import defaultdict
from contextlib import suppress
from functools import partial
from traceback import format_exception_only

import numpy as np
import pandas as pd
from scipy.stats import rankdata
from sklearn.base import clone
from sklearn.model_selection import ParameterGrid, ParameterSampler
from sklearn.model_selection import check_cv
from sklearn.model_selection._search import _check_param_grid
from sklearn.model_selection._validation import _aggregate_score_dicts
from sklearn.utils.metaestimators import if_delegate_has_method

from sktime.exceptions import FitFailedWarning
from sktime.exceptions import NotFittedError
from sktime.forecasting.base import BaseForecaster
from sktime.forecasting.base._base import DEFAULT_ALPHA
from sktime.utils.validation.forecasting import check_scoring
from sktime.utils.validation.forecasting import check_y_X


def _score(y_test, y_pred, scorer):
    """Evaluate forecasts"""
    if not isinstance(y_pred, pd.Series):
        raise NotImplementedError(
            "multi-step forecasting horizons with multiple cutoffs/windows "
            "are not supported yet"
        )

    # select only test points for which we have made predictions
    if not np.all(np.isin(y_pred.index, y_test.index)):
        raise IndexError("Predicted time points are not in test set")
    y_test = y_test.loc[y_pred.index]

    scores = {name: func(y_test, y_pred) for name, func in scorer.items()}
    return _check_scores(scores, scorer)


def _check_scores(scores, scorer):
    """Check returned scores"""
    error_msg = "scoring must return a number, got %s (%s) instead. (scorer=%s)"
    if isinstance(scores, dict):
        for name, score in scores.items():
            if hasattr(score, "item"):
                with suppress(ValueError):
                    # e.g. unwrap memmapped scalars
                    score = score.item()
            if not isinstance(score, numbers.Number):
                raise ValueError(error_msg % (score, type(score), name))
            scores[name] = score
    else:  # scalar
        if hasattr(scores, "item"):
            with suppress(ValueError):
                # e.g. unwrap memmapped scalars
                scores = scores.item()
        if not isinstance(scores, numbers.Number):
            raise ValueError(error_msg % (scores, type(scores), scorer))
    return scores


def _update_score(forecaster, cv, y, X, scorer):
    """Make, update and evaluate forecasts"""
    y_pred = forecaster.update_predict(y, cv=cv, X=X)
    return _score(y, y_pred, scorer)


def _split(y, X, cv):
    """Split data into training and validation window"""
    training_window, validation_window = cv.split_initial(y)
    y_train = y.iloc[training_window]
    y_val = y.iloc[validation_window]

    if X is not None:
        X_train = X.iloc[training_window, :]
        X_val = X.iloc[validation_window, :]
    else:
        X_train = None
        X_val = None

    return y_train, y_val, X_train, X_val


def _fit_and_score(
    forecaster,
    cv,
    y,
    X,
    scorer,
    verbose,
    parameters,
    fit_params,
    return_parameters=False,
    return_times=False,
    return_train_score=False,
    return_forecaster=False,
    error_score=np.nan,
):
    if return_train_score:
        raise NotImplementedError()

    # Get forecasting horizon
    fh = cv.get_fh()

    # Fit params
    fit_params = fit_params if fit_params is not None else {}
    if parameters is not None:
        forecaster.set_params(**parameters)

    # Split training data into training set and validation set
    y_train, y_val, X_train, X_val = _split(y, X, cv)

    # Fit forecaster
    start_time = time.time()
    try:
        forecaster.fit(y_train, X_train, fh)

    except Exception as e:
        # Note fit time as time until error
        fit_time = time.time() - start_time
        score_time = 0.0
        if error_score == "raise":
            raise
        elif isinstance(error_score, numbers.Number):
            if isinstance(scorer, dict):
                test_scores = {name: error_score for name in scorer}
            else:
                test_scores = error_score
            warnings.warn(
                "forecaster fit failed. The score on this train-test"
                " partition for these parameters will be set to %f. "
                "Details: \n%s" % (error_score, format_exception_only(type(e), e)[0]),
                FitFailedWarning,
            )
        else:
            raise ValueError(
                "error_score must be the string 'raise' or a"
                " numeric value. (Hint: if using 'raise', please"
                " make sure that it has been spelled correctly.)"
            )

    else:
        fit_time = time.time() - start_time
        test_scores = _update_score(forecaster, cv, y_val, X_val, scorer)
        score_time = time.time() - start_time - fit_time

    ret = [test_scores]

    if return_times:
        ret.extend([fit_time, score_time])
    if return_parameters:
        ret.append(parameters)
    if return_forecaster:
        ret.append(forecaster)
    return ret


class BaseGridSearch(BaseForecaster):
    def __init__(
        self,
        forecaster,
        cv,
        n_jobs=None,
        pre_dispatch=None,
        refit=False,
        scoring=None,
        verbose=0,
        error_score=None,
        return_train_score=None,
    ):
        self.forecaster = forecaster
        self.cv = cv
        self.n_jobs = n_jobs
        self.pre_dispatch = pre_dispatch
        self.refit = refit
        self.scoring = scoring
        self.verbose = verbose
        self.error_score = error_score
        self.return_train_score = return_train_score
        super(BaseGridSearch, self).__init__()

    @if_delegate_has_method(delegate=("best_forecaster_", "forecaster"))
    def update(self, y, X=None, update_params=False):
        """Call predict 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 update_predict(
        self,
        y,
        cv=None,
        X=None,
        update_params=False,
        return_pred_int=False,
        alpha=DEFAULT_ALPHA,
    ):
        """Call update_predict on the forecaster with the best found
        parameters.
        """
        self.check_is_fitted("update_predict")
        return self.best_forecaster_.update_predict(
            y,
            cv=cv,
            X=X,
            update_params=update_params,
            return_pred_int=return_pred_int,
            alpha=alpha,
        )

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

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

    @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):
        """Returns the score on the given data, if the forecaster has been
        refit.
        This uses the score defined by ``scoring`` where provided, and the
        ``best_forecaster_.score`` method otherwise.
        Parameters
        ----------
        X : pandas.DataFrame, shape=[n_obs, n_vars], optional (default=None)
            An optional 2-d dataframe of exogenous variables.
        y : pandas.Series
            Target time series to which to compare the forecasts.
        Returns
        -------
        score : float
        """
        self.check_is_fitted("score")
        if self.scorer_ is None:
            raise ValueError(
                "No score function explicitly defined, "
                "and the forecaster doesn't provide one %s" % self.best_forecaster_
            )
        score = self.scorer_
        y_pred = self.best_forecaster_.predict(fh, X)
        return score(y, y_pred)

    def _run_search(self, evaluate_candidates):
        raise NotImplementedError("_run_search not implemented.")

    @staticmethod
    def _format_results(candidate_params, scorers, out):
        n_candidates = len(candidate_params)
        (test_score_dicts, fit_time, score_time) = zip(*out)
        test_scores = _aggregate_score_dicts(test_score_dicts)

        results = {}

        def _store(key_name, array, rank=False, greater_is_better=False):
            """A small helper to store the scores/times to the cv_results_"""
            # When iterated first by splits, then by parameters
            # We want `array` to have `n_candidates` rows and `n_splits` cols.
            array = np.array(array, dtype=np.float64)

            results["mean_%s" % key_name] = array

            if rank:
                array = -array if greater_is_better else array
                results["rank_%s" % key_name] = np.asarray(
                    rankdata(array, method="min"), dtype=np.int32
                )

        _store("fit_time", fit_time)
        _store("score_time", score_time)
        # Use one MaskedArray and mask all the places where the param is not
        # applicable for that candidate. Use defaultdict as each candidate may
        # not contain all the params
        param_results = defaultdict(
            partial(
                np.ma.MaskedArray,
                np.empty(
                    n_candidates,
                ),
                mask=True,
                dtype=object,
            )
        )
        for cand_i, params in enumerate(candidate_params):
            for name, value in params.items():
                # An all masked empty array gets created for the key
                # `"param_%s" % name` at the first occurrence of `name`.
                # Setting the value at an index also unmasks that index
                param_results["param_%s" % name][cand_i] = value

        results.update(param_results)
        # Store a list of param dicts at the key "params"
        results["params"] = candidate_params

        for scorer_name, scorer in scorers.items():
            # Computed the (weighted) mean and std for test scores alone
            _store(
                "test_%s" % scorer_name,
                test_scores[scorer_name],
                rank=True,
                greater_is_better=scorer.greater_is_better,
            )

        return results

    def check_is_fitted(self, method_name=None):
        super(BaseGridSearch, self).check_is_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_params):
        """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.
        """
        y, X = check_y_X(y, X)

        # validate cross-validator
        cv = check_cv(self.cv)
        base_forecaster = clone(self.forecaster)

        scoring = check_scoring(self.scoring)
        scorers = {scoring.name: scoring}
        refit_metric = scoring.name

        fit_and_score_kwargs = dict(
            scorer=scorers,
            fit_params=fit_params,
            return_train_score=self.return_train_score,
            return_times=True,
            return_parameters=False,
            error_score=self.error_score,
            verbose=self.verbose,
        )

        results = {}
        all_candidate_params = []
        all_out = []

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

            if self.verbose > 0:
                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 = []
            for parameters in candidate_params:
                r = _fit_and_score(
                    clone(base_forecaster),
                    cv,
                    y,
                    X,
                    parameters=parameters,
                    **fit_and_score_kwargs
                )
                out.append(r)

            n_splits = cv.get_n_splits(y)

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

            all_candidate_params.extend(candidate_params)
            all_out.extend(out)

            nonlocal results
            results = self._format_results(all_candidate_params, scorers, all_out)
            return results

        self._run_search(evaluate_candidates)

        self.best_index_ = results["rank_test_%s" % refit_metric].argmin()
        self.best_score_ = results["mean_test_%s" % refit_metric][self.best_index_]
        self.best_params_ = results["params"][self.best_index_]

        self.best_forecaster_ = clone(base_forecaster).set_params(**self.best_params_)

        if self.refit:
            refit_start_time = time.time()
            self.best_forecaster_.fit(y, X, fh)
            self.refit_time_ = time.time() - refit_start_time

        # Store the only scorer not as a dict for single metric evaluation
        self.scorer_ = scorers[scoring.name]

        self.cv_results_ = results
        self.n_splits_ = cv.get_n_splits(y)

        self._is_fitted = True
        return self


class ForecastingGridSearchCV(BaseGridSearch):
    """
    Performs 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)
    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)

    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
    """

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

    def __init__(
        self,
        forecaster,
        cv,
        param_grid,
        scoring=None,
        n_jobs=None,
        refit=True,
        verbose=0,
        pre_dispatch="2*n_jobs",
        error_score=np.nan,
        return_train_score=False,
    ):
        super(ForecastingGridSearchCV, self).__init__(
            forecaster=forecaster,
            scoring=scoring,
            n_jobs=n_jobs,
            refit=refit,
            cv=cv,
            verbose=verbose,
            pre_dispatch=pre_dispatch,
            error_score=error_score,
            return_train_score=return_train_score,
        )
        self.param_grid = param_grid
        _check_param_grid(param_grid)

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


class ForecastingRandomizedSearchCV(BaseGridSearch):
    """
    Performs 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)
    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')
    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)

    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
    """

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

    def __init__(
        self,
        forecaster,
        cv,
        param_distributions,
        n_iter=10,
        scoring=None,
        n_jobs=None,
        refit=True,
        verbose=0,
        random_state=None,
        pre_dispatch="2*n_jobs",
        error_score=np.nan,
        return_train_score=False,
    ):
        super(ForecastingRandomizedSearchCV, self).__init__(
            forecaster=forecaster,
            scoring=scoring,
            n_jobs=n_jobs,
            refit=refit,
            cv=cv,
            verbose=verbose,
            pre_dispatch=pre_dispatch,
            error_score=error_score,
            return_train_score=return_train_score,
        )
        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"""
        evaluate_candidates(
            ParameterSampler(
                self.param_distributions, self.n_iter, random_state=self.random_state
            )
        )
