################################################################################
#                  skforecast.model_selection._validation                      #
#                                                                              #
# This work by skforecast team is licensed under the BSD 3-Clause License.     #
################################################################################
# coding=utf-8

import re
from copy import deepcopy
from typing import Union, Tuple, Optional, Callable
import warnings
import numpy as np
import pandas as pd
from joblib import Parallel, delayed, cpu_count
from tqdm.auto import tqdm
from ..metrics import add_y_train_argument, _get_metric
from ..exceptions import LongTrainingWarning, IgnoredArgumentWarning
from ..model_selection._split import TimeSeriesFold
from ..model_selection._utils import (
    _initialize_levels_model_selection_multiseries,
    check_backtesting_input,
    select_n_jobs_backtesting,
    _extract_data_folds_multiseries,
    _calculate_metrics_backtesting_multiseries
)
from ..utils import set_skforecast_warnings


def _backtesting_forecaster(
    forecaster: object,
    y: pd.Series,
    metric: Union[str, Callable, list],
    cv: TimeSeriesFold,
    exog: Optional[Union[pd.Series, pd.DataFrame]] = None,
    interval: Optional[list] = None,
    n_boot: int = 250,
    random_state: int = 123,
    use_in_sample_residuals: bool = True,
    use_binned_residuals: bool = False,
    n_jobs: Union[int, str] = 'auto',
    verbose: bool = False,
    show_progress: bool = True
) -> Tuple[pd.DataFrame, pd.DataFrame]:
    """
    Backtesting of forecaster model following the folds generated by the TimeSeriesFold
    class and using the metric(s) provided.

    If `forecaster` is already trained and `initial_train_size` is set to `None` in the
    TimeSeriesFold class, no initial train will be done and all data will be used
    to evaluate the model. However, the first `len(forecaster.last_window)` observations
    are needed to create the initial predictors, so no predictions are calculated for
    them.
    
    A copy of the original forecaster is created so that it is not modified during the
    process.
    
    Parameters
    ----------
    forecaster : ForecasterRecursive, ForecasterDirect
        Forecaster model.
    y : pandas Series
        Training time series.
    metric : str, Callable, list
        Metric used to quantify the goodness of fit of the model.
        
        - If `string`: {'mean_squared_error', 'mean_absolute_error',
        'mean_absolute_percentage_error', 'mean_squared_log_error',
        'mean_absolute_scaled_error', 'root_mean_squared_scaled_error'}
        - If `Callable`: Function with arguments `y_true`, `y_pred` and `y_train`
        (Optional) that returns a float.
        - If `list`: List containing multiple strings and/or Callables.
    cv : TimeSeriesFold
        TimeSeriesFold object with the information needed to split the data into folds.
        **New in version 0.14.0**
    exog : pandas Series, pandas DataFrame, default `None`
        Exogenous variable/s included as predictor/s. Must have the same
        number of observations as `y` and should be aligned so that y[i] is
        regressed on exog[i].
    interval : list, default `None`
        Confidence of the prediction interval estimated. Sequence of percentiles
        to compute, which must be between 0 and 100 inclusive. For example, 
        interval of 95% should be as `interval = [2.5, 97.5]`. If `None`, no
        intervals are estimated.
    n_boot : int, default `250`
        Number of bootstrapping iterations used to estimate prediction
        intervals.
    random_state : int, default `123`
        Sets a seed to the random generator, so that boot intervals are always 
        deterministic.
    use_in_sample_residuals : bool, default `True`
        If `True`, residuals from the training data are used as proxy of prediction 
        error to create prediction intervals. If `False`, out_sample_residuals 
        are used if they are already stored inside the forecaster.
    use_binned_residuals : bool, default `False`
        If `True`, residuals used in each bootstrapping iteration are selected
        conditioning on the predicted values. If `False`, residuals are selected
        randomly without conditioning on the predicted values.
    n_jobs : int, 'auto', default `'auto'`
        The number of jobs to run in parallel. If `-1`, then the number of jobs is 
        set to the number of cores. If 'auto', `n_jobs` is set using the function
        skforecast.utils.select_n_jobs_backtesting.
    verbose : bool, default `False`
        Print number of folds and index of training and validation sets used 
        for backtesting.
    show_progress : bool, default `True`
        Whether to show a progress bar.

    Returns
    -------
    metric_values : pandas DataFrame
        Value(s) of the metric(s).
    backtest_predictions : pandas DataFrame
        Value of predictions and their estimated interval if `interval` is not `None`.

        - column pred: predictions.
        - column lower_bound: lower bound of the interval.
        - column upper_bound: upper bound of the interval.
    
    """

    forecaster = deepcopy(forecaster)
    cv = deepcopy(cv)

    cv.set_params({
        'window_size': forecaster.window_size,
        'differentiation': forecaster.differentiation,
        'return_all_indexes': False,
        'verbose': verbose
    })

    initial_train_size = cv.initial_train_size
    refit = cv.refit
    gap = cv.gap
    
    if n_jobs == 'auto':
        n_jobs = select_n_jobs_backtesting(
                     forecaster = forecaster,
                     refit      = refit
                 )
    elif not isinstance(refit, bool) and refit != 1 and n_jobs != 1:
        warnings.warn(
            "If `refit` is an integer other than 1 (intermittent refit). `n_jobs` "
            "is set to 1 to avoid unexpected results during parallelization.",
            IgnoredArgumentWarning
        )
        n_jobs = 1
    else:
        n_jobs = n_jobs if n_jobs > 0 else cpu_count()

    if not isinstance(metric, list):
        metrics = [
            _get_metric(metric=metric)
            if isinstance(metric, str)
            else add_y_train_argument(metric)
        ]
    else:
        metrics = [
            _get_metric(metric=m)
            if isinstance(m, str)
            else add_y_train_argument(m) 
            for m in metric
        ]

    store_in_sample_residuals = False if interval is None else True

    folds = cv.split(X=y, as_pandas=False)
    window_size = cv.window_size

    if initial_train_size is not None:
        # First model training, this is done to allow parallelization when `refit`
        # is `False`. The initial Forecaster fit is outside the auxiliary function.
        exog_train = exog.iloc[:initial_train_size, ] if exog is not None else None
        forecaster.fit(
            y                         = y.iloc[:initial_train_size, ],
            exog                      = exog_train,
            store_in_sample_residuals = store_in_sample_residuals
        )
        # This is done to allow parallelization when `refit` is `False`. The initial 
        # Forecaster fit is outside the auxiliary function.
        folds[0][4] = False

    if refit:
        n_of_fits = int(len(folds) / refit)
        if type(forecaster).__name__ != 'ForecasterDirect' and n_of_fits > 50:
            warnings.warn(
                (f"The forecaster will be fit {n_of_fits} times. This can take substantial"
                 f" amounts of time. If not feasible, try with `refit = False`.\n"),
                LongTrainingWarning
            )
        elif type(forecaster).__name__ == 'ForecasterDirect' and n_of_fits * forecaster.steps > 50:
            warnings.warn(
                (f"The forecaster will be fit {n_of_fits * forecaster.steps} times "
                 f"({n_of_fits} folds * {forecaster.steps} regressors). This can take "
                 f"substantial amounts of time. If not feasible, try with `refit = False`.\n"),
                LongTrainingWarning
            )

    if show_progress:
        folds = tqdm(folds)

    def _fit_predict_forecaster(y, exog, forecaster, interval, fold, gap):
        """
        Fit the forecaster and predict `steps` ahead. This is an auxiliary 
        function used to parallelize the backtesting_forecaster function.
        """

        train_iloc_start       = fold[0][0]
        train_iloc_end         = fold[0][1]
        last_window_iloc_start = fold[1][0]
        last_window_iloc_end   = fold[1][1]
        test_iloc_start        = fold[2][0]
        test_iloc_end          = fold[2][1]

        if fold[4] is False:
            # When the model is not fitted, last_window must be updated to include
            # the data needed to make predictions.
            last_window_y = y.iloc[last_window_iloc_start:last_window_iloc_end]
        else:
            # The model is fitted before making predictions. If `fixed_train_size`
            # the train size doesn't increase but moves by `steps` in each iteration.
            # If `False` the train size increases by `steps` in each iteration.
            y_train = y.iloc[train_iloc_start:train_iloc_end, ]
            exog_train = (
                exog.iloc[train_iloc_start:train_iloc_end,] if exog is not None else None
            )
            last_window_y = None
            forecaster.fit(
                y                         = y_train, 
                exog                      = exog_train, 
                store_in_sample_residuals = store_in_sample_residuals
            )

        next_window_exog = exog.iloc[test_iloc_start:test_iloc_end, ] if exog is not None else None

        steps = len(range(test_iloc_start, test_iloc_end))
        if type(forecaster).__name__ == 'ForecasterDirect' and gap > 0:
            # Select only the steps that need to be predicted if gap > 0
            test_no_gap_iloc_start = fold[3][0]
            test_no_gap_iloc_end   = fold[3][1]
            steps = list(
                np.arange(len(range(test_no_gap_iloc_start, test_no_gap_iloc_end)))
                + gap
                + 1
            )

        if interval is None:
            pred = forecaster.predict(
                       steps       = steps,
                       last_window = last_window_y,
                       exog        = next_window_exog
                   )
        else:
            pred = forecaster.predict_interval(
                       steps                   = steps,
                       last_window             = last_window_y,
                       exog                    = next_window_exog,
                       interval                = interval,
                       n_boot                  = n_boot,
                       random_state            = random_state,
                       use_in_sample_residuals = use_in_sample_residuals,
                       use_binned_residuals    = use_binned_residuals,
                   )

        if type(forecaster).__name__ != 'ForecasterDirect' and gap > 0:
            pred = pred.iloc[gap:, ]

        return pred

    backtest_predictions = (
        Parallel(n_jobs=n_jobs)
        (delayed(_fit_predict_forecaster)
        (y=y, exog=exog, forecaster=forecaster, interval=interval, fold=fold, gap=gap)
         for fold in folds)
    )

    backtest_predictions = pd.concat(backtest_predictions)
    if isinstance(backtest_predictions, pd.Series):
        backtest_predictions = pd.DataFrame(backtest_predictions)

    train_indexes = []
    for i, fold in enumerate(folds):
        fit_fold = fold[-1]
        if i == 0 or fit_fold:
            train_iloc_start = fold[0][0] + window_size  # Exclude observations used to create predictors
            train_iloc_end = fold[0][1]
            train_indexes.append(np.arange(train_iloc_start, train_iloc_end))
    
    train_indexes = np.unique(np.concatenate(train_indexes))
    y_train = y.iloc[train_indexes]

    metric_values = [
        m(
            y_true = y.loc[backtest_predictions.index],
            y_pred = backtest_predictions['pred'],
            y_train = y_train
        ) 
        for m in metrics
    ]

    metric_values = pd.DataFrame(
        data    = [metric_values],
        columns = [m.__name__ for m in metrics]
    )
    
    return metric_values, backtest_predictions


def backtesting_forecaster(
    forecaster: object,
    y: pd.Series,
    cv: TimeSeriesFold,
    metric: Union[str, Callable, list],
    exog: Optional[Union[pd.Series, pd.DataFrame]] = None,
    interval: Optional[list] = None,
    n_boot: int = 250,
    random_state: int = 123,
    use_in_sample_residuals: bool = True,
    use_binned_residuals: bool = False,
    n_jobs: Union[int, str] = 'auto',
    verbose: bool = False,
    show_progress: bool = True
) -> Tuple[pd.DataFrame, pd.DataFrame]:
    """
    Backtesting of forecaster model following the folds generated by the TimeSeriesFold
    class and using the metric(s) provided.

    If `forecaster` is already trained and `initial_train_size` is set to `None` in the
    TimeSeriesFold class, no initial train will be done and all data will be used
    to evaluate the model. However, the first `len(forecaster.last_window)` observations
    are needed to create the initial predictors, so no predictions are calculated for
    them.
    
    A copy of the original forecaster is created so that it is not modified during 
    the process.

    Parameters
    ----------
    forecaster : ForecasterRecursive, ForecasterDirect
        Forecaster model.
    y : pandas Series
        Training time series.
    cv : TimeSeriesFold
        TimeSeriesFold object with the information needed to split the data into folds.
        **New in version 0.14.0**
    metric : str, Callable, list
        Metric used to quantify the goodness of fit of the model.
        
        - If `string`: {'mean_squared_error', 'mean_absolute_error',
        'mean_absolute_percentage_error', 'mean_squared_log_error',
        'mean_absolute_scaled_error', 'root_mean_squared_scaled_error'}
        - If `Callable`: Function with arguments `y_true`, `y_pred` and `y_train`
        (Optional) that returns a float.
        - If `list`: List containing multiple strings and/or Callables.
    exog : pandas Series, pandas DataFrame, default `None`
        Exogenous variable/s included as predictor/s. Must have the same
        number of observations as `y` and should be aligned so that y[i] is
        regressed on exog[i].
    interval : list, default `None`
        Confidence of the prediction interval estimated. Sequence of percentiles
        to compute, which must be between 0 and 100 inclusive. For example, 
        interval of 95% should be as `interval = [2.5, 97.5]`. If `None`, no
        intervals are estimated.
    n_boot : int, default `250`
        Number of bootstrapping iterations used to estimate prediction
        intervals.
    random_state : int, default `123`
        Sets a seed to the random generator, so that boot intervals are always 
        deterministic.
    use_in_sample_residuals : bool, default `True`
        If `True`, residuals from the training data are used as proxy of prediction 
        error to create prediction intervals. If `False`, out_sample_residuals 
        are used if they are already stored inside the forecaster.
    use_binned_residuals : bool, default `False`
        If `True`, residuals used in each bootstrapping iteration are selected
        conditioning on the predicted values. If `False`, residuals are selected
        randomly without conditioning on the predicted values.
    n_jobs : int, 'auto', default `'auto'`
        The number of jobs to run in parallel. If `-1`, then the number of jobs is 
        set to the number of cores. If 'auto', `n_jobs` is set using the function
        skforecast.utils.select_n_jobs_backtesting.
    verbose : bool, default `False`
        Print number of folds and index of training and validation sets used 
        for backtesting.
    show_progress : bool, default `True`
        Whether to show a progress bar.

    Returns
    -------
    metric_values : pandas DataFrame
        Value(s) of the metric(s).
    backtest_predictions : pandas DataFrame
        Value of predictions and their estimated interval if `interval` is not `None`.

        - column pred: predictions.
        - column lower_bound: lower bound of the interval.
        - column upper_bound: upper bound of the interval.
    
    """

    forecaters_allowed = [
        'ForecasterRecursive', 
        'ForecasterDirect',
        'ForecasterEquivalentDate'
    ]
    
    if type(forecaster).__name__ not in forecaters_allowed:
        raise TypeError(
            (f"`forecaster` must be of type {forecaters_allowed}, for all other types of "
             f" forecasters use the functions available in the other `model_selection` "
             f"modules.")
        )
    
    check_backtesting_input(
        forecaster              = forecaster,
        cv                      = cv,
        y                       = y,
        metric                  = metric,
        interval                = interval,
        n_boot                  = n_boot,
        random_state            = random_state,
        use_in_sample_residuals = use_in_sample_residuals,
        use_binned_residuals    = use_binned_residuals,
        n_jobs                  = n_jobs,
        show_progress           = show_progress
    )
    
    if type(forecaster).__name__ == 'ForecasterDirect' and \
       forecaster.steps < cv.steps + cv.gap:
        raise ValueError(
            (f"When using a ForecasterDirect, the combination of steps "
             f"+ gap ({cv.steps + cv.gap}) cannot be greater than the `steps` parameter "
             f"declared when the forecaster is initialized ({forecaster.steps}).")
        )
    
    metric_values, backtest_predictions = _backtesting_forecaster(
        forecaster              = forecaster,
        y                       = y,
        cv                      = cv,
        metric                  = metric,
        exog                    = exog,
        interval                = interval,
        n_boot                  = n_boot,
        random_state            = random_state,
        use_in_sample_residuals = use_in_sample_residuals,
        use_binned_residuals    = use_binned_residuals,
        n_jobs                  = n_jobs,
        verbose                 = verbose,
        show_progress           = show_progress
    )

    return metric_values, backtest_predictions


def _backtesting_forecaster_multiseries(
    forecaster: object,
    series: Union[pd.DataFrame, dict],
    cv: TimeSeriesFold,
    metric: Union[str, Callable, list],
    levels: Optional[Union[str, list]] = None,
    add_aggregated_metric: bool = True,
    exog: Optional[Union[pd.Series, pd.DataFrame, dict]] = None,
    interval: Optional[list] = None,
    n_boot: int = 250,
    random_state: int = 123,
    use_in_sample_residuals: bool = True,
    n_jobs: Union[int, str] = 'auto',
    verbose: bool = False,
    show_progress: bool = True,
    suppress_warnings: bool = False
) -> Tuple[pd.DataFrame, pd.DataFrame]:
    """
    Backtesting of forecaster model following the folds generated by the TimeSeriesFold
    class and using the metric(s) provided.

    If `forecaster` is already trained and `initial_train_size` is set to `None` in the
    TimeSeriesFold class, no initial train will be done and all data will be used
    to evaluate the model. However, the first `len(forecaster.last_window)` observations
    are needed to create the initial predictors, so no predictions are calculated for
    them.
    
    A copy of the original forecaster is created so that it is not modified during the
    process.
    
    Parameters
    ----------
    forecaster : ForecasterRecursiveMultiSeries, ForecasterDirectMultiVariate
        Forecaster model.
    series : pandas DataFrame, dict
        Training time series.
    cv : TimeSeriesFold
        TimeSeriesFold object with the information needed to split the data into folds.
        **New in version 0.14.0**
    metric : str, Callable, list
        Metric used to quantify the goodness of fit of the model.
        
        - If `string`: {'mean_squared_error', 'mean_absolute_error',
        'mean_absolute_percentage_error', 'mean_squared_log_error'}
        - If `Callable`: Function with arguments y_true, y_pred that returns a float.
        - If `list`: List containing multiple strings and/or Callables.
    levels : str, list, default `None`
        Time series to be predicted. If `None` all levels will be predicted.
    add_aggregated_metric : bool, default `False`
        If `True`, and multiple series (`levels`) are predicted, the aggregated
        metrics (average, weighted average and pooled) are also returned.

        - 'average': the average (arithmetic mean) of all levels.
        - 'weighted_average': the average of the metrics weighted by the number of
        predicted values of each level.
        - 'pooling': the values of all levels are pooled and then the metric is
        calculated.
    exog : pandas Series, pandas DataFrame, dict, default `None`
        Exogenous variables.
    interval : list, default `None`
        Confidence of the prediction interval estimated. Sequence of percentiles
        to compute, which must be between 0 and 100 inclusive. If `None`, no
        intervals are estimated.
    n_boot : int, default `250`
        Number of bootstrapping iterations used to estimate prediction
        intervals.
    random_state : int, default `123`
        Sets a seed to the random generator, so that boot intervals are always 
        deterministic.
    use_in_sample_residuals : bool, default `True`
        If `True`, residuals from the training data are used as proxy of prediction
        error to create prediction intervals. If `False`, out_sample_residuals 
        are used if they are already stored inside the forecaster.
    n_jobs : int, 'auto', default `'auto'`
        The number of jobs to run in parallel. If `-1`, then the number of jobs is 
        set to the number of cores. If 'auto', `n_jobs` is set using the function
        skforecast.utils.select_n_jobs_backtesting.
    verbose : bool, default `False`
        Print number of folds and index of training and validation sets used 
        for backtesting.
    show_progress : bool, default `True`
        Whether to show a progress bar.
    suppress_warnings: bool, default `False`
        If `True`, skforecast warnings will be suppressed during the backtesting 
        process. See skforecast.exceptions.warn_skforecast_categories for more
        information.

    Returns
    -------
    metrics_levels : pandas DataFrame
        Value(s) of the metric(s). Index are the levels and columns the metrics.
    backtest_predictions : pandas Dataframe
        Value of predictions and their estimated interval if `interval` is not `None`. 
        If there is more than one level, this structure will be repeated for each of them.

        - column pred: predictions.
        - column lower_bound: lower bound of the interval.
        - column upper_bound: upper bound of the interval.
    
    """

    set_skforecast_warnings(suppress_warnings, action='ignore')

    forecaster = deepcopy(forecaster)
    cv = deepcopy(cv)

    cv.set_params({
        'window_size': forecaster.window_size,
        'differentiation': forecaster.differentiation,
        'return_all_indexes': False,
        'verbose': verbose
    })

    initial_train_size = cv.initial_train_size
    refit = cv.refit
    gap = cv.gap

    if n_jobs == 'auto':
        n_jobs = select_n_jobs_backtesting(
                     forecaster = forecaster,
                     refit      = refit
                 )
    elif not isinstance(refit, bool) and refit != 1 and n_jobs != 1:
        warnings.warn(
            ("If `refit` is an integer other than 1 (intermittent refit). `n_jobs` "
             "is set to 1 to avoid unexpected results during parallelization."),
             IgnoredArgumentWarning
        )
        n_jobs = 1
    else:
        n_jobs = n_jobs if n_jobs > 0 else cpu_count()

    levels = _initialize_levels_model_selection_multiseries(
                 forecaster = forecaster,
                 series     = series,
                 levels     = levels
             )

    if not isinstance(metric, list):
        metrics = [
            _get_metric(metric=metric)
            if isinstance(metric, str)
            else add_y_train_argument(metric)
        ]
    else:
        metrics = [
            _get_metric(metric=m)
            if isinstance(m, str)
            else add_y_train_argument(m) 
            for m in metric
        ]

    store_in_sample_residuals = False if interval is None else True

    folds = cv.split(X=series, as_pandas=False)
    span_index = cv._extract_index(X=series)

    if initial_train_size is not None:
        # First model training, this is done to allow parallelization when `refit`
        # is `False`. The initial Forecaster fit is outside the auxiliary function.
        data_fold = _extract_data_folds_multiseries(
                        series             = series,
                        folds              = [folds[0]],
                        span_index         = span_index,
                        window_size        = forecaster.window_size,
                        exog               = exog,
                        dropna_last_window = forecaster.dropna_from_series,
                        externally_fitted  = False
                    )
        series_train, _, last_window_levels, exog_train, _, _ = next(data_fold)
        forecaster.fit(
            series                    = series_train,
            exog                      = exog_train,
            store_last_window         = last_window_levels,
            store_in_sample_residuals = store_in_sample_residuals,
            suppress_warnings         = suppress_warnings
        )
        # This is done to allow parallelization when `refit` is `False`. The initial 
        # Forecaster fit is outside the auxiliary function.
        folds[0][4] = False
        
    if refit:
        n_of_fits = int(len(folds) / refit)
        if type(forecaster).__name__ != 'ForecasterDirectMultiVariate' and n_of_fits > 50:
            warnings.warn(
                (f"The forecaster will be fit {n_of_fits} times. This can take substantial "
                 f"amounts of time. If not feasible, try with `refit = False`.\n"),
                LongTrainingWarning,
            )
        elif type(forecaster).__name__ == 'ForecasterDirectMultiVariate' and n_of_fits * forecaster.steps > 50:
            warnings.warn(
                (f"The forecaster will be fit {n_of_fits * forecaster.steps} times "
                 f"({n_of_fits} folds * {forecaster.steps} regressors). This can take "
                 f"substantial amounts of time. If not feasible, try with `refit = False`.\n"),
                LongTrainingWarning
            )

    if show_progress:
        folds = tqdm(folds)
        
    externally_fitted = True if initial_train_size is None else False
    data_folds = _extract_data_folds_multiseries(
                     series             = series,
                     folds              = folds,
                     span_index         = span_index,
                     window_size        = forecaster.window_size,
                     exog               = exog,
                     dropna_last_window = forecaster.dropna_from_series,
                     externally_fitted  = externally_fitted
                 )

    def _fit_predict_forecaster(data_fold, forecaster, interval, levels, gap):
        """
        Fit the forecaster and predict `steps` ahead. This is an auxiliary 
        function used to parallelize the backtesting_forecaster_multiseries
        function.
        """

        (
            series_train,
            last_window_series,
            last_window_levels,
            exog_train,
            next_window_exog,
            fold
        ) = data_fold

        if fold[4] is True:
            forecaster.fit(
                series                    = series_train, 
                exog                      = exog_train,
                store_last_window         = last_window_levels,
                store_in_sample_residuals = store_in_sample_residuals,
                suppress_warnings         = suppress_warnings
            )

        test_iloc_start = fold[2][0]
        test_iloc_end   = fold[2][1]
        steps = len(range(test_iloc_start, test_iloc_end))
        if type(forecaster).__name__ == 'ForecasterDirectMultiVariate' and gap > 0:
            # Select only the steps that need to be predicted if gap > 0
            test_iloc_start = fold[3][0]
            test_iloc_end   = fold[3][1]
            steps = list(np.arange(len(range(test_iloc_start, test_iloc_end))) + gap + 1)

        levels_predict = [level for level in levels 
                          if level in last_window_levels]
        if interval is None:

            pred = forecaster.predict(
                       steps             = steps, 
                       levels            = levels_predict, 
                       last_window       = last_window_series,
                       exog              = next_window_exog,
                       suppress_warnings = suppress_warnings
                   )
        else:
            pred = forecaster.predict_interval(
                       steps                   = steps,
                       levels                  = levels_predict, 
                       last_window             = last_window_series,
                       exog                    = next_window_exog,
                       interval                = interval,
                       n_boot                  = n_boot,
                       random_state            = random_state,
                       use_in_sample_residuals = use_in_sample_residuals,
                       suppress_warnings       = suppress_warnings
                   )

        if type(forecaster).__name__ != 'ForecasterDirectMultiVariate' and gap > 0:
            pred = pred.iloc[gap:, ]

        return pred

    backtest_predictions = Parallel(n_jobs=n_jobs)(
        delayed(_fit_predict_forecaster)(
            data_fold  = data_fold,
            forecaster = forecaster,
            interval   = interval,
            levels     = levels,
            gap        = gap
        )
        for data_fold in data_folds
    )

    backtest_predictions = pd.concat(backtest_predictions, axis=0)

    levels_in_backtest_predictions = backtest_predictions.columns
    if interval is not None:
        levels_in_backtest_predictions = [
            level 
            for level in levels_in_backtest_predictions
            if not re.search(r'_lower_bound|_upper_bound', level)
        ]
    for level in levels_in_backtest_predictions:
        valid_index = series[level][series[level].notna()].index
        no_valid_index = backtest_predictions.index.difference(valid_index, sort=False)
        cols = [level]
        if interval:
            cols = cols + [f'{level}_lower_bound', f'{level}_upper_bound']
        backtest_predictions.loc[no_valid_index, cols] = np.nan

    metrics_levels = _calculate_metrics_backtesting_multiseries(
        series                = series,
        predictions           = backtest_predictions,
        folds                 = folds,
        span_index            = span_index,
        window_size           = forecaster.window_size,
        metrics               = metrics,
        levels                = levels,
        add_aggregated_metric = add_aggregated_metric
    )
       
    set_skforecast_warnings(suppress_warnings, action='default')

    return metrics_levels, backtest_predictions


def backtesting_forecaster_multiseries(
    forecaster: object,
    series: Union[pd.DataFrame, dict],
    cv: TimeSeriesFold,
    metric: Union[str, Callable, list],
    levels: Optional[Union[str, list]] = None,
    add_aggregated_metric: bool = True,
    exog: Optional[Union[pd.Series, pd.DataFrame, dict]] = None,
    interval: Optional[list] = None,
    n_boot: int = 250,
    random_state: int = 123,
    use_in_sample_residuals: bool = True,
    n_jobs: Union[int, str] = 'auto',
    verbose: bool = False,
    show_progress: bool = True,
    suppress_warnings: bool = False
) -> Tuple[pd.DataFrame, pd.DataFrame]:
    """
    Backtesting of forecaster model following the folds generated by the TimeSeriesFold
    class and using the metric(s) provided.

    If `forecaster` is already trained and `initial_train_size` is set to `None` in the
    TimeSeriesFold class, no initial train will be done and all data will be used
    to evaluate the model. However, the first `len(forecaster.last_window)` observations
    are needed to create the initial predictors, so no predictions are calculated for
    them.
    
    A copy of the original forecaster is created so that it is not modified during 
    the process.

    Parameters
    ----------
    forecaster : ForecasterRecursiveMultiSeries, ForecasterDirectMultiVariate, ForecasterRnn
        Forecaster model.
    series : pandas DataFrame, dict
        Training time series.
    cv : TimeSeriesFold
        TimeSeriesFold object with the information needed to split the data into folds.
    metric : str, Callable, list
        Metric used to quantify the goodness of fit of the model.
        
        - If `string`: {'mean_squared_error', 'mean_absolute_error',
        'mean_absolute_percentage_error', 'mean_squared_log_error',
        'mean_absolute_scaled_error', 'root_mean_squared_scaled_error'}
        - If `Callable`: Function with arguments `y_true`, `y_pred` and `y_train`
        (Optional) that returns a float.
        - If `list`: List containing multiple strings and/or Callables.
    levels : str, list, default `None`
        Time series to be predicted. If `None` all levels will be predicted.
    add_aggregated_metric : bool, default `True`
        If `True`, and multiple series (`levels`) are predicted, the aggregated
        metrics (average, weighted average and pooled) are also returned.

        - 'average': the average (arithmetic mean) of all levels.
        - 'weighted_average': the average of the metrics weighted by the number of
        predicted values of each level.
        - 'pooling': the values of all levels are pooled and then the metric is
        calculated.
    exog : pandas Series, pandas DataFrame, dict, default `None`
        Exogenous variables.
    interval : list, default `None`
        Confidence of the prediction interval estimated. Sequence of percentiles
        to compute, which must be between 0 and 100 inclusive. If `None`, no
        intervals are estimated.
    n_boot : int, default `250`
        Number of bootstrapping iterations used to estimate prediction
        intervals.
    random_state : int, default `123`
        Sets a seed to the random generator, so that boot intervals are always 
        deterministic.
    use_in_sample_residuals : bool, default `True`
        If `True`, residuals from the training data are used as proxy of prediction 
        error to create prediction intervals. If `False`, out_sample_residuals 
        are used if they are already stored inside the forecaster.
    n_jobs : int, 'auto', default `'auto'`
        The number of jobs to run in parallel. If `-1`, then the number of jobs is 
        set to the number of cores. If 'auto', `n_jobs` is set using the function
        skforecast.utils.select_n_jobs_backtesting.
    verbose : bool, default `False`
        Print number of folds and index of training and validation sets used 
        for backtesting.
    show_progress : bool, default `True`
        Whether to show a progress bar.
    suppress_warnings: bool, default `False`
        If `True`, skforecast warnings will be suppressed during the backtesting 
        process. See skforecast.exceptions.warn_skforecast_categories for more
        information.

    Returns
    -------
    metrics_levels : pandas DataFrame
        Value(s) of the metric(s). Index are the levels and columns the metrics.
    backtest_predictions : pandas DataFrame
        Value of predictions and their estimated interval if `interval` is not `None`.
        If there is more than one level, this structure will be repeated for each of them.

        - column pred: predictions.
        - column lower_bound: lower bound of the interval.
        - column upper_bound: upper bound of the interval.
    
    """

    multi_series_forecasters = [
        'ForecasterRecursiveMultiSeries', 
        'ForecasterDirectMultiVariate',
        'ForecasterRnn'
    ]

    forecaster_name = type(forecaster).__name__

    if forecaster_name not in multi_series_forecasters:
        raise TypeError(
            (f"`forecaster` must be of type {multi_series_forecasters}, "
             f"for all other types of forecasters use the functions available in "
             f"the `model_selection` module. Got {forecaster_name}")
        )
    
    check_backtesting_input(
        forecaster              = forecaster,
        cv                      = cv,
        metric                  = metric,
        add_aggregated_metric   = add_aggregated_metric,
        series                  = series,
        exog                    = exog,
        interval                = interval,
        n_boot                  = n_boot,
        random_state            = random_state,
        use_in_sample_residuals = use_in_sample_residuals,
        n_jobs                  = n_jobs,
        show_progress           = show_progress,
        suppress_warnings       = suppress_warnings
    )

    metrics_levels, backtest_predictions = _backtesting_forecaster_multiseries(
        forecaster              = forecaster,
        series                  = series,
        cv                      = cv,
        levels                  = levels,
        metric                  = metric,
        add_aggregated_metric   = add_aggregated_metric,
        exog                    = exog,
        interval                = interval,
        n_boot                  = n_boot,
        random_state            = random_state,
        use_in_sample_residuals = use_in_sample_residuals,
        n_jobs                  = n_jobs,
        verbose                 = verbose,
        show_progress           = show_progress,
        suppress_warnings       = suppress_warnings
    )

    return metrics_levels, backtest_predictions


def _backtesting_sarimax(
    forecaster: object,
    y: pd.Series,
    metric: Union[str, Callable, list],
    cv: TimeSeriesFold,
    exog: Optional[Union[pd.Series, pd.DataFrame]] = None,
    alpha: Optional[float] = None,
    interval: Optional[list] = None,
    n_jobs: Union[int, str] = 'auto',
    suppress_warnings_fit: bool = False,
    verbose: bool = False,
    show_progress: bool = True,
) -> Tuple[pd.DataFrame, pd.DataFrame]:
    """
    Backtesting of ForecasterSarimax.
    
    A copy of the original forecaster is created so that it is not modified during 
    the process.
    
    Parameters
    ----------
    forecaster : ForecasterSarimax
        Forecaster model.
    y : pandas Series
        Training time series.
    metric : str, Callable, list
        Metric used to quantify the goodness of fit of the model.
        
        - If `string`: {'mean_squared_error', 'mean_absolute_error',
        'mean_absolute_percentage_error', 'mean_squared_log_error',
        'mean_absolute_scaled_error', 'root_mean_squared_scaled_error'}
        - If `Callable`: Function with arguments `y_true`, `y_pred` and `y_train`
        (Optional) that returns a float.
        - If `list`: List containing multiple strings and/or Callables.
    cv : TimeSeriesFold
        TimeSeriesFold object with the information needed to split the data into folds.
        **New in version 0.14.0**
    exog : pandas Series, pandas DataFrame, default `None`
        Exogenous variable/s included as predictor/s. Must have the same
        number of observations as `y` and should be aligned so that y[i] is
        regressed on exog[i].
    alpha : float, default `0.05`
        The confidence intervals for the forecasts are (1 - alpha) %.
        If both, `alpha` and `interval` are provided, `alpha` will be used.
    interval : list, default `None`
        Confidence of the prediction interval estimated. The values must be
        symmetric. Sequence of percentiles to compute, which must be between 
        0 and 100 inclusive. For example, interval of 95% should be as 
        `interval = [2.5, 97.5]`. If both, `alpha` and `interval` are 
        provided, `alpha` will be used.
    n_jobs : int, 'auto', default `'auto'`
        The number of jobs to run in parallel. If `-1`, then the number of jobs is 
        set to the number of cores. If 'auto', `n_jobs` is set using the function
        skforecast.utils.select_n_jobs_backtesting.
    verbose : bool, default `False`
        Print number of folds and index of training and validation sets used 
        for backtesting.
    suppress_warnings_fit : bool, default `False`
        If `True`, warnings generated during fitting will be ignored.
    show_progress : bool, default `True`
        Whether to show a progress bar.

    Returns
    -------
    metric_values : pandas DataFrame
        Value(s) of the metric(s).
    backtest_predictions : pandas DataFrame
        Value of predictions and their estimated interval if `interval` is not `None`.

        - column pred: predictions.
        - column lower_bound: lower bound of the interval.
        - column upper_bound: upper bound of the interval.
    
    """

    forecaster = deepcopy(forecaster)
    cv = deepcopy(cv)

    cv.set_params({
        'window_size': forecaster.window_size,
        'return_all_indexes': False,
        'verbose': verbose
    })

    steps = cv.steps
    initial_train_size = cv.initial_train_size
    refit = cv.refit
    gap = cv.gap
    
    if refit == False:
        if n_jobs != 'auto' and n_jobs != 1:
            warnings.warn(
                ("If `refit = False`, `n_jobs` is set to 1 to avoid unexpected "
                 "results during parallelization."),
                 IgnoredArgumentWarning
            )
        n_jobs = 1
    else:
        if n_jobs == 'auto':        
            n_jobs = select_n_jobs_backtesting(
                         forecaster = forecaster,
                         refit      = refit
                     )
        elif not isinstance(refit, bool) and refit != 1 and n_jobs != 1:
            warnings.warn(
                ("If `refit` is an integer other than 1 (intermittent refit). `n_jobs` "
                 "is set to 1 to avoid unexpected results during parallelization."),
                 IgnoredArgumentWarning
            )
            n_jobs = 1
        else:
            n_jobs = n_jobs if n_jobs > 0 else cpu_count()

    if not isinstance(metric, list):
        metrics = [
            _get_metric(metric=metric)
            if isinstance(metric, str)
            else add_y_train_argument(metric)
        ]
    else:
        metrics = [
            _get_metric(metric=m)
            if isinstance(m, str)
            else add_y_train_argument(m) 
            for m in metric
        ]

    # initial_train_size cannot be None because of append method in Sarimax
    # First model training, this is done to allow parallelization when `refit` 
    # is `False`. The initial Forecaster fit is outside the auxiliary function.
    exog_train = exog.iloc[:initial_train_size, ] if exog is not None else None
    forecaster.fit(
        y                 = y.iloc[:initial_train_size, ],
        exog              = exog_train,
        suppress_warnings = suppress_warnings_fit
    )
    
    folds = cv.split(X=y, as_pandas=False)
    # This is done to allow parallelization when `refit` is `False`. The initial 
    # Forecaster fit is outside the auxiliary function.
    folds[0][4] = False
    
    if refit:
        n_of_fits = int(len(folds) / refit)
        if n_of_fits > 50:
            warnings.warn(
                (f"The forecaster will be fit {n_of_fits} times. This can take substantial "
                 f"amounts of time. If not feasible, try with `refit = False`.\n"),
                 LongTrainingWarning
            )
       
    folds_tqdm = tqdm(folds) if show_progress else folds

    def _fit_predict_forecaster(y, exog, forecaster, alpha, interval, fold, steps, gap):
        """
        Fit the forecaster and predict `steps` ahead. This is an auxiliary 
        function used to parallelize the backtesting_forecaster function.
        """

        # In each iteration the model is fitted before making predictions. 
        # if fixed_train_size the train size doesn't increase but moves by `steps` 
        # in each iteration. if False the train size increases by `steps` in each 
        # iteration.
        train_idx_start = fold[0][0]
        train_idx_end   = fold[0][1]
        test_idx_start  = fold[2][0]
        test_idx_end    = fold[2][1]

        if refit:
            last_window_start = fold[0][1]  # Same as train_idx_end
            last_window_end   = fold[1][1]
        else:
            last_window_end   = fold[2][0]  # test_idx_start
            last_window_start = last_window_end - steps

        if fold[4] is False:
            # When the model is not fitted, last_window and last_window_exog must 
            # be updated to include the data needed to make predictions.
            last_window_y = y.iloc[last_window_start:last_window_end]
            last_window_exog = exog.iloc[last_window_start:last_window_end] if exog is not None else None 
        else:
            # The model is fitted before making predictions. If `fixed_train_size`  
            # the train size doesn't increase but moves by `steps` in each iteration. 
            # If `False` the train size increases by `steps` in each  iteration.
            y_train = y.iloc[train_idx_start:train_idx_end, ]
            exog_train = exog.iloc[train_idx_start:train_idx_end, ] if exog is not None else None
            
            last_window_y = None
            last_window_exog = None

            forecaster.fit(y=y_train, exog=exog_train, suppress_warnings=suppress_warnings_fit)

        next_window_exog = exog.iloc[test_idx_start:test_idx_end, ] if exog is not None else None

        # After the first fit, Sarimax must use the last windows stored in the model
        if fold == folds[0]:
            last_window_y = None
            last_window_exog = None

        steps = len(range(test_idx_start, test_idx_end))
        if alpha is None and interval is None:
            pred = forecaster.predict(
                       steps            = steps,
                       last_window      = last_window_y,
                       last_window_exog = last_window_exog,
                       exog             = next_window_exog
                   )
        else:
            pred = forecaster.predict_interval(
                       steps            = steps,
                       exog             = next_window_exog,
                       alpha            = alpha,
                       interval         = interval,
                       last_window      = last_window_y,
                       last_window_exog = last_window_exog
                   )

        pred = pred.iloc[gap:, ]            
        
        return pred

    backtest_predictions = (
        Parallel(n_jobs=n_jobs)
        (delayed(_fit_predict_forecaster)
        (
            y=y, 
            exog=exog, 
            forecaster=forecaster, 
            alpha=alpha, 
            interval=interval, 
            fold=fold, 
            steps=steps,
            gap=gap
        )
        for fold in folds_tqdm)
    )
    
    backtest_predictions = pd.concat(backtest_predictions)
    if isinstance(backtest_predictions, pd.Series):
        backtest_predictions = pd.DataFrame(backtest_predictions)

    train_indexes = []
    for i, fold in enumerate(folds):
        fit_fold = fold[-1]
        if i == 0 or fit_fold:
            train_iloc_start = fold[0][0]
            train_iloc_end = fold[0][1]
            train_indexes.append(np.arange(train_iloc_start, train_iloc_end))
    
    train_indexes = np.unique(np.concatenate(train_indexes))
    y_train = y.iloc[train_indexes]

    metric_values = [
        m(
            y_true = y.loc[backtest_predictions.index],
            y_pred = backtest_predictions['pred'],
            y_train = y_train
        ) 
        for m in metrics
    ]

    metric_values = pd.DataFrame(
        data    = [metric_values],
        columns = [m.__name__ for m in metrics]
    )

    return metric_values, backtest_predictions


def backtesting_sarimax(
    forecaster: object,
    y: pd.Series,
    cv: TimeSeriesFold,
    metric: Union[str, Callable, list],
    exog: Optional[Union[pd.Series, pd.DataFrame]] = None,
    alpha: Optional[float] = None,
    interval: Optional[list] = None,
    n_jobs: Union[int, str] = 'auto',
    verbose: bool = False,
    suppress_warnings_fit: bool = False,
    show_progress: bool = True
) -> Tuple[pd.DataFrame, pd.DataFrame]:
    """
    Backtesting of ForecasterSarimax.
    
    A copy of the original forecaster is created so that it is not modified during 
    the process.

    Parameters
    ----------
    forecaster : ForecasterSarimax
        Forecaster model.
    y : pandas Series
        Training time series.
    cv : TimeSeriesFold
        TimeSeriesFold object with the information needed to split the data into folds.
        **New in version 0.14.0**
    metric : str, Callable, list
        Metric used to quantify the goodness of fit of the model.
        
        - If `string`: {'mean_squared_error', 'mean_absolute_error',
        'mean_absolute_percentage_error', 'mean_squared_log_error',
        'mean_absolute_scaled_error', 'root_mean_squared_scaled_error'}
        - If `Callable`: Function with arguments `y_true`, `y_pred` and `y_train`
        (Optional) that returns a float.
        - If `list`: List containing multiple strings and/or Callables.
    exog : pandas Series, pandas DataFrame, default `None`
        Exogenous variable/s included as predictor/s. Must have the same
        number of observations as `y` and should be aligned so that y[i] is
        regressed on exog[i].
    alpha : float, default `0.05`
        The confidence intervals for the forecasts are (1 - alpha) %.
        If both, `alpha` and `interval` are provided, `alpha` will be used.
    interval : list, default `None`
        Confidence of the prediction interval estimated. The values must be
        symmetric. Sequence of percentiles to compute, which must be between 
        0 and 100 inclusive. For example, interval of 95% should be as 
        `interval = [2.5, 97.5]`. If both, `alpha` and `interval` are 
        provided, `alpha` will be used.
    n_jobs : int, 'auto', default `'auto'`
        The number of jobs to run in parallel. If `-1`, then the number of jobs is 
        set to the number of cores. If 'auto', `n_jobs` is set using the function
        skforecast.utils.select_n_jobs_backtesting. 
    verbose : bool, default `False`
        Print number of folds and index of training and validation sets used 
        for backtesting.
    suppress_warnings_fit : bool, default `False`
        If `True`, warnings generated during fitting will be ignored.
    show_progress : bool, default `True`
        Whether to show a progress bar.

    Returns
    -------
    metric_values : pandas DataFrame
        Value(s) of the metric(s).
    backtest_predictions : pandas DataFrame
        Value of predictions and their estimated interval if `interval` is not `None`.

        - column pred: predictions.
        - column lower_bound: lower bound of the interval.
        - column upper_bound: upper bound of the interval.
    
    """
    
    if type(forecaster).__name__ not in ['ForecasterSarimax']:
        raise TypeError(
            ("`forecaster` must be of type `ForecasterSarimax`, for all other "
             "types of forecasters use the functions available in the other "
             "`model_selection` modules.")
        )
    
    check_backtesting_input(
        forecaster            = forecaster,
        cv                    = cv,
        y                     = y,
        metric                = metric,
        interval              = interval,
        alpha                 = alpha,
        n_jobs                = n_jobs,
        show_progress         = show_progress,
        suppress_warnings_fit = suppress_warnings_fit
    )
    
    metric_values, backtest_predictions = _backtesting_sarimax(
        forecaster            = forecaster,
        y                     = y,
        cv                    = cv,
        metric                = metric,
        exog                  = exog,
        alpha                 = alpha,
        interval              = interval,
        n_jobs                = n_jobs,
        verbose               = verbose,
        suppress_warnings_fit = suppress_warnings_fit,
        show_progress         = show_progress
    )

    return metric_values, backtest_predictions
