# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/experiments__auto.ipynb (unless otherwise specified).

__all__ = ['auto', 'instantiate_space']

# Cell
import numpy as np
import pandas as pd
import neuralforecast as nf
from hyperopt import hp

# Cell
def auto(config_dict: dict, Y_df: pd.DataFrame, X_df: pd.DataFrame, S_df: pd.DataFrame,
           loss_function_val: callable, loss_functions_test: dict, forecast_horizon: int, ts_in_val: int, ts_in_test: int,
           return_forecasts: bool = False, return_model: bool = True, test_auto: bool = False, verbose: bool = False):
    """
    Auto hyperparameter tuning function.

    Parameters
    ----------
    config_dict: Dict
        Dictionary with configuration. Keys should be name of models.
        For each model specify the hyperparameter space (None will use default suggested space), hyperopt steps and timeout.
    Y_df: pd.DataFrame
        Target time series with columns ['unique_id', 'ds', 'y'].
    X_df: pd.DataFrame
        Exogenous time series with columns ['unique_id', 'ds', 'y'].
    S_df: pd.DataFrame
        Static exogenous variables with columns ['unique_id', 'ds'].
        and static variables.
    loss_function_val: function
        Loss function used for validation.
    loss_functions_test: Dictionary
        Loss functions used for test.
        (function name: string, function: fun)
    forecast_horizon: int
        Forecast horizon
    ts_in_val: int
        Number of timestamps in validation.
    ts_in_test: int
        Number of timestamps in test.
    return_forecasts: bool
        If true return forecast on test.
    return_model: bool
        If true return model.
    test_auto: bool
        If true, will only run one training step and hyperopt iteration for each model. Use it for testing purposes.
    verbose:
        If true, will print summary of dataset, model and training.
    """

    if test_auto:
        print('WARNING: THIS IS A TEST, MODELS WILL NOT BE TRAINED PROPERLY!')

    # Data characteristics
    n_series = Y_df['unique_id'].nunique()
    n_x = len(X_df.columns)-2 if X_df is not None else 0
    n_s = len(S_df.columns)-1 if S_df is not None else 0
    frequency = pd.infer_freq(Y_df['ds'])

    # Hyperopt
    models = config_dict.keys()
    output_dict = {}
    best_model = None
    best_loss = np.inf
    for model in models:
        model_config = config_dict[model]

        hyperopt_steps = model_config['hyperopt_steps'] if (test_auto==False) else 1

        if model_config['space'] is None:
            space = instantiate_space(model=model, n_time_out=forecast_horizon, n_series=n_series, n_x=n_x, n_s=n_s, frequency=frequency, test=test_auto)
        else:
            space = model_config['space']
        trials = nf.experiments.utils.hyperopt_tunning(space=space,
                                                       hyperopt_max_evals=hyperopt_steps,
                                                       loss_function_val=loss_function_val, loss_functions_test=loss_functions_test,
                                                       S_df=S_df, Y_df=Y_df, X_df=X_df, f_cols=[], ds_in_val=ts_in_val, ds_in_test=ts_in_test,
                                                       return_forecasts=return_forecasts,
                                                       return_model=return_model,
                                                       save_progress=False,
                                                       results_file=None,
                                                       step_save_progress=0,
                                                       verbose=verbose)

        model_output = {'best_mc': trials.best_trial['result']['mc'],
                        'run_time': trials.best_trial['result']['run_time'],
                        'best_val_loss': trials.best_trial['result']['loss']
                       }

        # Return model
        if return_model:
            model_output['model'] = trials.best_trial['result']['model']

        # Return test losses
        if ts_in_test > 0:
            model_output['best_test_loss'] = trials.best_trial['result']['test_losses']

        # Return test forecasts
        if (return_forecasts) and (ts_in_test > 0):
            model_output['y_hat'] = trials.best_trial['result']['forecasts_test']['test_y_hat']
            model_output['y_true'] = trials.best_trial['result']['forecasts_test']['test_y_true']

        # Improvement
        optimization_times = [trials.trials[0]['result']['loss']]
        optimization_losses = [trials.trials[0]['result']['run_time']]
        for i in range(1, len(trials)):
            loss = trials.trials[i]['result']['loss']
            time = trials.trials[i]['result']['run_time']

            if loss > np.min(optimization_losses):
                loss = np.min(optimization_losses)
            optimization_losses.append(loss)
            optimization_times.append(np.sum(optimization_times)+time)

        model_output['optimization_losses'] = optimization_losses
        model_output['optimization_times'] = optimization_times

        # Append to dict
        output_dict[model] = model_output

        if trials.best_trial['result']['loss'] < best_loss:
            best_model = trials.best_trial['result']['model']
            best_loss = trials.best_trial['result']['loss']

    return best_model, output_dict

# Cell
def instantiate_space(model, n_time_out, n_series, n_x, n_s, frequency, test):
    assert model in ['nbeats', 'nhits', 'rnn'], f'Invalid model {model}'

    if model == 'nbeats':
        space = nf.models.nbeats.nbeats.suggested_space(n_time_out=n_time_out, n_series=n_series, n_x=n_x, n_s=n_s, frequency=frequency)
    elif model == 'nhits':
        space = nf.models.nhits.nhits.suggested_space(n_time_out=n_time_out, n_series=n_series, n_x=n_x, n_s=n_s, frequency=frequency)
    elif model == 'rnn':
        space = nf.models.rnn.rnn.suggested_space(n_time_out=n_time_out, n_series=n_series, n_x=n_x, n_s=n_s, frequency=frequency)

    if test:
        space['max_steps'] = hp.choice('max_steps', [1])
        space['max_epochs'] = hp.choice('max_epochs', [None])

    return space