content: import os
  import platform
  import shutil
  import sys
  import tempfile
  import warnings
  from datetime import date
  from pathlib import Path
  
  import git
  import numpy as np
  import optuna
  import pandas as pd
  import polars
  import polars.testing
  import pytest
  import s3fs
  import torch
  from ray import tune
  
  from neuralforecast.auto import (
      MLP,
      NHITS,
      RNN,
      TCN,
      TFT,
      AutoDilatedRNN,
      Autoformer,
      AutoMLP,
      AutoNBEATSx,
      AutoRNN,
      DeepAR,
      DilatedRNN,
      Informer,
      NBEATSx,
      StemGNN,
      TSMixer,
      TSMixerx,
      VanillaTransformer,
  )
  from neuralforecast.common.enums import ExplainerEnum
  from neuralforecast.core import (
      LSTM,
      DLinear,
      FEDformer,
      NeuralForecast,
      PatchTST,
      PredictionIntervals,
      TimesNet,
      _insample_times,
      _type2scaler,
  )
  from neuralforecast.losses.pytorch import (
      GMM,
      MAE,
      NBMM,
      PMM,
      DistributionLoss,
      MQLoss,
  )
  from neuralforecast.tsdataset import TimeSeriesDataset
  from neuralforecast.utils import (
      AirPassengersPanel,
      AirPassengersStatic,
      generate_series,
  )
  
  from .test_helpers import assert_equal_dfs, get_expected_size
  
  
  @pytest.fixture
  def setup():
      uids = pd.Series(["id_0", "id_1"])
      indptr = np.array([0, 4, 10], dtype=np.int32)
      return uids, indptr
  
  
  @pytest.mark.parametrize("step_size, freq, days", [(1, "D", 1), (2, "W-THU", 14)])
  def test_cutoff_deltas(setup, step_size, freq, days):
      uids, indptr = setup
      h = 2
      times = np.hstack(
              [
                  pd.date_range("2000-01-01", freq=freq, periods=4),
                  pd.date_range("2000-10-10", freq=freq, periods=10),
              ]
          )
      times_df = _insample_times(times, uids, indptr, h, freq, step_size=step_size)
      pd.testing.assert_frame_equal(
          times_df.groupby("unique_id")["ds"].min().reset_index(),
          pd.DataFrame(
              {
                  "unique_id": uids,
                  "ds": times[indptr[:-1]],
              }
          ),
      )
      pd.testing.assert_frame_equal(
          times_df.groupby("unique_id")["ds"].max().reset_index(),
          pd.DataFrame(
              {
                  "unique_id": uids,
                  "ds": times[indptr[1:] - 1],
              }
          ),
      )
      cutoff_deltas = (
          times_df.drop_duplicates(["unique_id", "cutoff"])
          .groupby("unique_id")["cutoff"]
          .diff()
          .dropna()
      )
      assert cutoff_deltas.nunique() == 1
      assert cutoff_deltas.unique()[0] == pd.Timedelta(f"{days}D")
  
  @pytest.fixture
  def setup_airplane_data_polars(setup_airplane_data):
      """Create polars version of airplane data with renamed columns."""
      AirPassengersPanel_train, AirPassengersPanel_test = setup_airplane_data
  
      # Set up column renaming for Polars
      renamer = {"unique_id": "uid", "ds": "time", "y": "target"}
  
      # Create Polars dataframes
      AirPassengers_pl = polars.from_pandas(AirPassengersPanel_train)
      AirPassengers_pl = AirPassengers_pl.rename(renamer)
  
      AirPassengersStatic_pl = polars.from_pandas(AirPassengersStatic)
      AirPassengersStatic_pl = AirPassengersStatic_pl.rename({'unique_id': 'uid'})
  
      return AirPassengers_pl, AirPassengersStatic_pl
  
  
  # Unittest for early stopping without val_size protection
  def test_neural_forecast_early_stopping(setup_airplane_data):
      AirPassengersPanel_train, _ = setup_airplane_data
      models = [NHITS(h=12, input_size=12, max_steps=1, early_stop_patience_steps=5)]
      nf = NeuralForecast(models=models, freq="M")
      with pytest.raises(Exception, match="Set val_size>0 if early stopping is enabled."):
          nf.fit(AirPassengersPanel_train)
  
  
  # test fit+cross_validation behaviour
  def test_neural_forecast_fit_cross_validation(setup_airplane_data):
      AirPassengersPanel_train, _ = setup_airplane_data
      models = [NHITS(h=12, input_size=24, max_steps=10)]
      nf = NeuralForecast(models=models, freq="M")
      nf.fit(AirPassengersPanel_train)
      init_fcst = nf.predict()
      init_cv = nf.cross_validation(AirPassengersPanel_train, use_init_models=True)
      after_cv = nf.cross_validation(AirPassengersPanel_train, use_init_models=True)
      nf.fit(AirPassengersPanel_train, use_init_models=True)
      after_fcst = nf.predict()
      pd.testing.assert_frame_equal(init_cv, after_cv)
      pd.testing.assert_frame_equal(after_fcst, init_fcst)
  
  # test cross_validation with refit
  def test_neural_forecast_refit(setup_airplane_data):
      AirPassengersPanel_train, _ = setup_airplane_data
      models = [
          NHITS(
              h=12,
              input_size=24,
              max_steps=2,
              futr_exog_list=["trend"],
              stat_exog_list=["airline1", "airline2"],
          )
      ]
      nf = NeuralForecast(models=models, freq="M")
      cv_kwargs = dict(
          df=AirPassengersPanel_train,
          static_df=AirPassengersStatic,
          n_windows=4,
          use_init_models=True,
      )
      cv_res_norefit = nf.cross_validation(refit=False, **cv_kwargs)
      cutoffs = cv_res_norefit["cutoff"].unique()
      for refit in [True, 2]:
          cv_res = nf.cross_validation(refit=refit, **cv_kwargs)
          refit = int(refit)
          fltr = lambda df: df["cutoff"].isin(cutoffs[:refit])
          expected = cv_res_norefit[fltr]
          actual = cv_res[fltr]
          # predictions for the no-refit windows should be the same
          pd.testing.assert_frame_equal(
              actual.reset_index(drop=True), expected.reset_index(drop=True)
          )
          # predictions after refit should be different
          with pytest.raises(AssertionError, match=r'\(column name="NHITS"\) are different'):
              pd.testing.assert_frame_equal(
                  cv_res_norefit.drop(expected.index).reset_index(drop=True),
                  cv_res.drop(actual.index).reset_index(drop=True),
              )
  
  
  def test_neural_forecast_scaling(setup_airplane_data):
      """Test scaling functionality for NeuralForecast models."""
      AirPassengersPanel_train, AirPassengersPanel_test = setup_airplane_data
  
      # Get initial forecast for comparison
      init_models = [NHITS(h=12, input_size=24, max_steps=10)]
      nf_init = NeuralForecast(models=init_models, freq="M")
      nf_init.fit(AirPassengersPanel_train)
      init_fcst = nf_init.predict()
  
      models = [NHITS(h=12, input_size=24, max_steps=10)]
      models_exog = [
          NHITS(
              h=12,
              input_size=12,
              max_steps=10,
              hist_exog_list=["trend"],
              futr_exog_list=["trend"],
          )
      ]
  
      # Test fit+predict with standard scaling
      nf = NeuralForecast(models=models, freq="M", local_scaler_type="standard")
      nf.fit(AirPassengersPanel_train)
      scaled_fcst = nf.predict()
      # Check that the forecasts are similar to the one without scaling
      np.testing.assert_allclose(
          init_fcst["NHITS"].values,
          scaled_fcst["NHITS"].values,
          rtol=0.3,
      )
  
      # Test with exogenous variables
      nf = NeuralForecast(models=models_exog, freq="M", local_scaler_type="standard")
      nf.fit(AirPassengersPanel_train)
      scaled_exog_fcst = nf.predict(futr_df=AirPassengersPanel_test)
      # Check that the forecasts are similar to the one without exog
      np.testing.assert_allclose(
          scaled_fcst["NHITS"].values,
          scaled_exog_fcst["NHITS"].values,
          rtol=0.3,
      )
  
      # Test cross-validation with robust scaling
      nf = NeuralForecast(models=models, freq="M", local_scaler_type="robust")
      cv_res = nf.cross_validation(AirPassengersPanel)
      # Check that the forecasts are similar to the original values
      np.testing.assert_allclose(
          cv_res["NHITS"].values,
          cv_res["y"].values,
          rtol=0.3,
      )
  
      # Test cross-validation with exogenous variables and robust-iqr scaling
      nf = NeuralForecast(models=models_exog, freq="M", local_scaler_type="robust-iqr")
      cv_res_exog = nf.cross_validation(AirPassengersPanel)
      # Check that the forecasts are similar to the original values
      np.testing.assert_allclose(
          cv_res_exog["NHITS"].values,
          cv_res_exog["y"].values,
          rtol=0.2,
      )
  
      # Test fit+predict_insample with minmax scaling
      nf = NeuralForecast(models=models, freq="M", local_scaler_type="minmax")
      nf.fit(AirPassengersPanel_train)
      insample_res = (
          nf.predict_insample()
          .groupby("unique_id")
          .tail(-12)  # first values aren't reliable
          .merge(
              AirPassengersPanel_train[["unique_id", "ds", "y"]],
              on=["unique_id", "ds"],
              how="left",
              suffixes=("_actual", "_expected"),
          )
      )
      # Check that y is inverted correctly
      np.testing.assert_allclose(
          insample_res["y_actual"].values,
          insample_res["y_expected"].values,
          rtol=1e-5,
      )
      # Check that predictions are in the same scale
      np.testing.assert_allclose(
          insample_res["NHITS"].values,
          insample_res["y_expected"].values,
          rtol=0.7,
      )
  
      # Test with exogenous variables
      nf = NeuralForecast(models=models_exog, freq="M", local_scaler_type="minmax")
      nf.fit(AirPassengersPanel_train)
      insample_res_exog = (
          nf.predict_insample()
          .groupby("unique_id")
          .tail(-12)  # first values aren't reliable
          .merge(
              AirPassengersPanel_train[["unique_id", "ds", "y"]],
              on=["unique_id", "ds"],
              how="left",
              suffixes=("_actual", "_expected"),
          )
      )
      # Check that y is inverted correctly
      np.testing.assert_allclose(
          insample_res_exog["y_actual"].values,
          insample_res_exog["y_expected"].values,
          rtol=1e-5,
      )
      # Check that predictions are similar to those without exog
      np.testing.assert_allclose(
          insample_res["NHITS"].values,
          insample_res_exog["NHITS"].values,
          rtol=0.2,
      )
  
  
  def test_neural_forecast_boxcox_scaling(setup_airplane_data):
      """Test BoxCox scaling functionality for NeuralForecast models."""
      AirPassengersPanel_train, _ = setup_airplane_data
  
      models = [NHITS(h=12, input_size=24, max_steps=10)]
  
      # Test BoxCox scaling
      nf = NeuralForecast(models=models, freq="M", local_scaler_type="boxcox")
      nf.fit(AirPassengersPanel_train)
      insample_res = (
          nf.predict_insample()
          .groupby("unique_id")
          .tail(-12)  # first values aren't reliable
          .merge(
              AirPassengersPanel_train[["unique_id", "ds", "y"]],
              on=["unique_id", "ds"],
              how="left",
              suffixes=("_actual", "_expected"),
          )
      )
      # Check that y is inverted correctly
      np.testing.assert_allclose(
          insample_res["y_actual"].values,
          insample_res["y_expected"].values,
          rtol=1e-5,
      )
      # Check that predictions are in the same scale
      np.testing.assert_allclose(
          insample_res["NHITS"].values,
          insample_res["y_expected"].values,
          rtol=0.7,
      )
  
  
  # Test static exogenous feature scaling
  
  @pytest.fixture
  def config() -> dict:
      return {
          'h': 12,
          'input_size': 24,
          'max_steps': 10,
      }
  
  
  @pytest.mark.parametrize("scaler", _type2scaler.keys())
  def test_neural_forecast_static_scaling(config, scaler):
      """Test static scaling functionality for NeuralForecast models."""
      stat_exog = ['airline1', 'airline2']
      nf = NeuralForecast(models=[NHITS(**config, stat_exog_list=stat_exog)], freq="D")
      nf.fit(AirPassengersPanel, AirPassengersStatic)
      without_scaler = nf.predict()
  
      nf = NeuralForecast(models=[NHITS(**config, stat_exog_list=stat_exog)], freq="D", local_static_scaler_type=scaler)
      nf.fit(AirPassengersPanel, AirPassengersStatic)
      with_scaler = nf.predict()
  
      np.testing.assert_allclose(
          without_scaler["NHITS"].values,
          with_scaler["NHITS"].values,
          rtol=0.2,
      )
  
  
  @pytest.fixture
  def data(size=300, n_series=3) -> pd.DataFrame:
      return pd.DataFrame({
          "unique_id": (['store_1'] * (size // n_series) + ['store_2'] * (size // n_series) + ['store_3'] * (size // n_series)),
          "ds": np.tile(pd.date_range(start="2020-01-01", periods=size // n_series, freq="D").to_numpy(), n_series),
          "y": np.random.rand(size),
      })
  
  
  @pytest.fixture
  def static_data(n_series=3) -> pd.DataFrame:
      return pd.DataFrame({
          "unique_id": [f'store_{i+1}' for i in range(n_series)],
          "size": np.random.randint(50, 500, size=n_series),
          "num_days_open": np.random.randint(100, 1000, size=n_series),
      })
  
  
  @pytest.mark.parametrize("scaler", _type2scaler.keys())
  def test_normalization_of_static_exog(data, static_data, config, scaler):
      models = [TFT(**config, stat_exog_list=["size", "num_days_open"])]
      nf = NeuralForecast(models=models, freq="D", local_static_scaler_type=scaler)
      
      dataset = TimeSeriesDataset.from_df(data, static_data)[0]
      nf._scalers_fit_transform(dataset)
      fit_dataset = dataset.static
      
      dataset = TimeSeriesDataset.from_df(data, static_data)[0]
      nf._scalers_transform(dataset)
      predict_dataset = dataset.static
  
      assert (fit_dataset == predict_dataset).all()
  
  
  def test_standard_normalization_of_static_exog(data, static_data, config):
      models = [TFT(**config, stat_exog_list=["size", "num_days_open"])]
      nf = NeuralForecast(models=models, freq="D", local_static_scaler_type="standard")
      
      dataset = TimeSeriesDataset.from_df(data, static_data)[0]
      nf._scalers_fit_transform(dataset)
      normalized_static = dataset.static.numpy()
  
      for i, col in enumerate(dataset.static_cols):
          col_values = static_data[col].values.reshape(-1, 1).astype(np.float32)
          mean = col_values.mean()
          std = col_values.std()
          expected_normalized = (col_values - mean) / std
  
          np.testing.assert_allclose(nf.static_scalers_[col].stats_[:, 0], mean, rtol=1e-5)
          np.testing.assert_allclose(nf.static_scalers_[col].stats_[:, 1], std, rtol=1e-5)
          np.testing.assert_allclose(
              normalized_static[:, i],
              expected_normalized.flatten(),
              rtol=1e-5,
          )
  
  
  def test_minmax_normalization_of_static_exog(data, static_data, config):
      models = [TFT(**config, stat_exog_list=["size", "num_days_open"])]
      nf = NeuralForecast(models=models, freq="D", local_static_scaler_type="minmax")
      
      dataset = TimeSeriesDataset.from_df(data, static_data)[0]
      nf._scalers_fit_transform(dataset)
      normalized_static = dataset.static.numpy()
  
      for i, col in enumerate(dataset.static_cols):
          col_values = static_data[col].values.reshape(-1, 1).astype(np.float32)
          min_val = col_values.min()
          range_val = col_values.max() - min_val
          expected_normalized = (col_values - min_val) / range_val
  
          np.testing.assert_allclose(nf.static_scalers_[col].stats_[:, 0], min_val, rtol=1e-5)
          np.testing.assert_allclose(nf.static_scalers_[col].stats_[:, 1], range_val, rtol=1e-5)
          np.testing.assert_allclose(
              normalized_static[:, i],
              expected_normalized.flatten(),
              rtol=1e-5,
          )
  
  
  # test futr_df contents
  def test_future_df_contents(setup_airplane_data):
      AirPassengersPanel_train, AirPassengersPanel_test = setup_airplane_data
  
      models = [
          NHITS(
              h=6,
              input_size=24,
              max_steps=10,
              hist_exog_list=["trend"],
              futr_exog_list=["trend"],
          )
      ]
      nf = NeuralForecast(models=models, freq="M")
      nf.fit(AirPassengersPanel_train)
  
      # Test that not enough rows in futr_df raises an error
      with pytest.raises(Exception) as exc_info:
          nf.predict(futr_df=AirPassengersPanel_test.head())
      assert "There are missing combinations" in str(exc_info.value)
  
      # Test that extra rows issues a warning
      with warnings.catch_warnings(record=True) as issued_warnings:
          warnings.simplefilter("always", UserWarning)
          nf.predict(futr_df=AirPassengersPanel_test)
      assert any("Dropped 12 unused rows" in str(w.message) for w in issued_warnings)
  
      # Test that models require futr_df and not providing it raises an error
      with pytest.raises(Exception) as exc_info:
          nf.predict()
      assert "Models require the following future exogenous features: {'trend'}" in str(exc_info.value)
  
      # Test that missing feature in futr_df raises an error
      with pytest.raises(Exception) as exc_info:
          nf.predict(futr_df=AirPassengersPanel_test.drop(columns="trend"))
      assert "missing from `futr_df`: {'trend'}" in str(exc_info.value)
  
      # Test that null values in futr_df raises an error
      with pytest.raises(Exception) as exc_info:
          nf.predict(futr_df=AirPassengersPanel_test.assign(trend=np.nan))
      assert "Found null values in `futr_df`" in str(exc_info.value)
  
  # Test inplace model fitting
  def test_inplace_model_fitting(setup_airplane_data):
      AirPassengersPanel_train, AirPassengersPanel_test = setup_airplane_data
  
      models = [MLP(h=12, input_size=12, max_steps=1, scaler_type="robust")]
      initial_weights = models[0].mlp[0].weight.detach().clone()
      fcst = NeuralForecast(models=models, freq="M")
      fcst.fit(
          df=AirPassengersPanel_train, static_df=AirPassengersStatic, use_init_models=True
      )
      after_weights = fcst.models_init[0].mlp[0].weight.detach().clone()
      assert np.allclose(initial_weights, after_weights), "init models should not be modified"
      assert len(fcst.models[0].train_trajectories) > 0, (
          "models stored trajectories should not be empty"
      )
  
  @pytest.fixture
  def setup_models_for_insample():
      h = 12
  
      models = [
          NHITS(
              h=h,
              input_size=24,
              loss=MQLoss(level=[80]),
              max_steps=1,
              alias="NHITS",
              scaler_type=None,
          ),
          RNN(h=h, input_size=-1, loss=MAE(), max_steps=1, alias="RNN", scaler_type=None),
      ]
      return models
  
  
  
  # Test predict_insample
  def test_predict_insample(setup_airplane_data, setup_models_for_insample):
      AirPassengersPanel_train, AirPassengersPanel_test = setup_airplane_data
      models = setup_models_for_insample
      test_size = 12
      h = 12
  
      nf = NeuralForecast(models=models, freq="M")
      _ = nf.cross_validation(
          df=AirPassengersPanel_train,
          static_df=AirPassengersStatic,
          val_size=0,
          test_size=test_size,
          n_windows=None,
      )
  
      forecasts = nf.predict_insample(step_size=1)
  
      expected_size = get_expected_size(AirPassengersPanel_train, h, test_size, step_size=1)
      assert len(forecasts) == expected_size, (
          f"Shape mismatch in predict_insample: {len(forecasts)=}, {expected_size=}"
      )
  
  # Test predict_insample (different lengths)
  def test_predict_insample_diff_lengths(setup_models_for_insample):
      models = setup_models_for_insample
      test_size = 12
      n_series = 2
      h = 12
      diff_len_df = generate_series(n_series=n_series, max_length=100)
  
      nf = NeuralForecast(models=models, freq="D")
      _ = nf.cross_validation(
          df=diff_len_df, val_size=0, test_size=test_size, n_windows=None
      )
  
      forecasts = nf.predict_insample(step_size=1)
      expected_size = get_expected_size(diff_len_df, h, test_size, step_size=1)
      assert len(forecasts) == expected_size, (
          f"Shape mismatch in predict_insample: {len(forecasts)=}, {expected_size=}"
      )
  
  @pytest.mark.parametrize("step_size, test_size", [(7, 0), (9, 0), (7, 5), (9, 5)])
  def test_predict_insample_step_size(setup_airplane_data, step_size, test_size):
      AirPassengersPanel_train, _ = setup_airplane_data
  
      h = 12
      train_end = AirPassengersPanel_train['ds'].max()
      sizes = AirPassengersPanel_train['unique_id'].value_counts().to_numpy()
  
      models = [NHITS(h=h, input_size=12, max_steps=1)]
      nf = NeuralForecast(models=models, freq='M')
      nf.fit(AirPassengersPanel_train)
      # Note: only apply set_test_size() upon nf.fit(), otherwise it would have set the test_size = 0
      nf.models[0].set_test_size(test_size)
  
      forecasts = nf.predict_insample(step_size=step_size)
      last_cutoff = train_end - test_size * pd.offsets.MonthEnd() - h * pd.offsets.MonthEnd()
      n_expected_cutoffs = (sizes[0] - test_size - nf.h + step_size) // step_size
  
      # compare cutoff values
      expected_cutoffs = np.flip(np.array([last_cutoff - step_size * i * pd.offsets.MonthEnd() for i in range(n_expected_cutoffs)]))
      actual_cutoffs = np.array([pd.Timestamp(x) for x in forecasts[forecasts['unique_id']==nf.uids[1]]['cutoff'].unique()])
      np.testing.assert_array_equal(expected_cutoffs, actual_cutoffs, err_msg=f"{step_size=},{expected_cutoffs=},{actual_cutoffs=}")
  
      # check forecast-points count per series
      cutoffs_by_series = forecasts.groupby(['unique_id', 'cutoff']).size().unstack('unique_id')
      pd.testing.assert_series_equal(cutoffs_by_series['Airline1'], cutoffs_by_series['Airline2'], check_names=False)
  
  
  def test_predict_insample_diff_loss(setup_airplane_data):
      AirPassengersPanel_train, _ = setup_airplane_data
  
  
      def get_expected_cols(model, level):
          # index columns
          n_cols = 4
          for model in models:
              if isinstance(loss, (DistributionLoss, PMM, GMM, NBMM)):
                  if level is None:
                      # Variations of DistributionLoss return the sample mean as well
                      n_cols += len(loss.quantiles) + 1
                  else:
                      # Variations of DistributionLoss return the sample mean as well
                      n_cols += 2 * len(level) + 1
              else:
                  if level is None:
                      # Other probabilistic models return the sample mean as well
                      n_cols += 1
                  # Other probabilistic models return just the levels
                  else:
                      n_cols += len(level) + 1
          return n_cols
  
      for loss in [
          # IQLoss(),
          DistributionLoss(distribution="Normal", level=[80]),
          PMM(level=[80]),
      ]:
          for level in [None, [80, 90]]:
              # Use CPU accelerator on macOS to avoid MPS LSTM projection limitation
              # MPS (Metal Performance Shaders) doesn't support LSTM with projections
              # which are enabled when recurrent=True
              lstm_kwargs = {"h": 12, "input_size": 12, "loss": loss, "max_steps": 1, "recurrent": True}
              if platform.system() == "Darwin":  # macOS
                  lstm_kwargs["accelerator"] = "cpu"
  
              models = [
                  NHITS(h=12, input_size=12, loss=loss, max_steps=1),
                  LSTM(**lstm_kwargs),
              ]
              nf = NeuralForecast(models=models, freq='D')
  
              nf.fit(df=AirPassengersPanel_train)
              df = nf.predict_insample(step_size=1, level=level)
              expected_cols = get_expected_cols(models, level)
              assert df.shape[1] == expected_cols, f'Shape mismatch for {loss} and level={level} in predict_insample: cols={df.shape[1]}, expected_cols={expected_cols}'
  
  # Test aliases
  def test_aliases(setup_airplane_data):
      AirPassengersPanel_train, AirPassengersPanel_test = setup_airplane_data
  
      config_drnn = {
          "input_size": tune.choice([-1]),
          "encoder_hidden_size": tune.choice([5, 10]),
          "max_steps": 1,
          "val_check_steps": 1,
          "step_size": 1,
      }
      models = [
          # test Auto
          AutoDilatedRNN(h=12, config=config_drnn, cpus=1, num_samples=2, alias="AutoDIL"),
          # test BaseWindows
          NHITS(h=12, input_size=24, loss=MQLoss(level=[80]), max_steps=1, alias="NHITSMQ"),
          # test BaseRecurrent
          RNN(
              h=12,
              input_size=-1,
              encoder_hidden_size=10,
              max_steps=1,
              stat_exog_list=["airline1"],
              futr_exog_list=["trend"],
              hist_exog_list=["y_[lag12]"],
              alias="MyRNN",
          ),
          # test BaseMultivariate
          StemGNN(
              h=12,
              input_size=24,
              n_series=2,
              max_steps=1,
              scaler_type="robust",
              alias="StemMulti",
          ),
          # test model without alias
          NHITS(h=12, input_size=24, max_steps=1),
      ]
      nf = NeuralForecast(models=models, freq="M")
      nf.fit(df=AirPassengersPanel_train, static_df=AirPassengersStatic)
      forecasts = nf.predict(futr_df=AirPassengersPanel_test)
      assert forecasts.columns.to_list() == [
              "unique_id",
              "ds",
              "AutoDIL",
              "NHITSMQ-median",
              "NHITSMQ-lo-80",
              "NHITSMQ-hi-80",
              "MyRNN",
              "StemMulti",
              "NHITS",
          ]
  
  def config_optuna(trial):
      return {
          "input_size": trial.suggest_categorical("input_size", [12, 24]),
          "hist_exog_list": trial.suggest_categorical(
              "hist_exog_list", [["trend"], ["y_[lag12]"], ["trend", "y_[lag12]"]]
          ),
          "futr_exog_list": ["trend"],
          "max_steps": 10,
          "val_check_steps": 5,
      }
  
  def test_training_with_an_iterative_dataset(setup_airplane_data):
      AirPassengersPanel_train, AirPassengersPanel_test = setup_airplane_data
  
      config_ray = {
          "input_size": tune.choice([12, 24]),
          "hist_exog_list": tune.choice([["trend"], ["y_[lag12]"], ["trend", "y_[lag12]"]]),
          "futr_exog_list": ["trend"],
          "max_steps": 10,
          "val_check_steps": 5,
      }
      # test training with an iterative dataset produces the same results as directly passing in the dataset as a pandas dataframe
      AirPassengersPanel_train["id"] = AirPassengersPanel_train["unique_id"]
      AirPassengersPanel_test["id"] = AirPassengersPanel_test["unique_id"]
  
      models = [
          NHITS(h=12, input_size=12, max_steps=10, futr_exog_list=["trend"], random_seed=1),
          AutoMLP(
              h=12,
              config=config_optuna,
              num_samples=2,
              backend="optuna",
              search_alg=optuna.samplers.TPESampler(seed=0),
          ),  # type: ignore
          AutoNBEATSx(h=12, config=config_ray, cpus=1, num_samples=2),
      ]
      nf = NeuralForecast(models=models, freq="M")
  
      # fit+predict with pandas dataframe
      nf.fit(
          df=AirPassengersPanel_train.drop(columns="unique_id"),
          use_init_models=True,
          id_col="id",
      )
      pred_dataframe = nf.predict(
          futr_df=AirPassengersPanel_test.drop(columns="unique_id")
      ).reset_index()
  
      # fit+predict with data directory
      with tempfile.TemporaryDirectory() as tmpdir:
          AirPassengersPanel_train.to_parquet(
              tmpdir, partition_cols=["unique_id"], index=False
          )
          data_directory = sorted([str(path) for path in Path(tmpdir).iterdir()])
          nf.fit(df=data_directory, use_init_models=True, id_col="id")
  
      pred_df = AirPassengersPanel_train[
          AirPassengersPanel_train["unique_id"] == "Airline2"
      ].drop(columns="unique_id")
      futr_df = AirPassengersPanel_test[
          AirPassengersPanel_test["unique_id"] == "Airline2"
      ].drop(columns="unique_id")
  
      pred_iterative = nf.predict(df=pred_df, futr_df=futr_df)
      pred_airline2 = pred_dataframe[pred_dataframe["id"] == "Airline2"]
      np.testing.assert_allclose(
          pred_iterative["NHITS"], pred_airline2["NHITS"], rtol=0, atol=1
      )
      np.testing.assert_allclose(
          pred_iterative["AutoMLP"], pred_airline2["AutoMLP"], rtol=0, atol=1
      )
      np.testing.assert_allclose(
          pred_iterative["AutoNBEATSx"], pred_airline2["AutoNBEATSx"], rtol=0, atol=1
      )
  
  
  # test cross validation no leakage
  def test_cross_validation(h=12, test_size=12):
      df = AirPassengersPanel
      static_df = AirPassengersStatic
      if (test_size - h) % 1:
          raise Exception("`test_size - h` should be module `step_size`")
  
      Y_test_df = df.groupby("unique_id").tail(test_size)
      Y_train_df = df.drop(Y_test_df.index)
      config = {
          "input_size": tune.choice([12, 24]),
          "step_size": 12,
          "hidden_size": 256,
          "max_steps": 1,
          "val_check_steps": 1,
      }
      config_drnn = {
          "input_size": tune.choice([-1]),
          "encoder_hidden_size": tune.choice([5, 10]),
          "max_steps": 1,
          "val_check_steps": 1,
      }
      fcst = NeuralForecast(
          models=[
              AutoDilatedRNN(h=12, config=config_drnn, cpus=1, num_samples=1),
              DilatedRNN(h=12, input_size=-1, encoder_hidden_size=5, max_steps=1),
              RNN(
                  h=12,
                  input_size=-1,
                  encoder_hidden_size=5,
                  max_steps=1,
                  stat_exog_list=["airline1"],
                  futr_exog_list=["trend"],
                  hist_exog_list=["y_[lag12]"],
              ),
              TCN(
                  h=12,
                  input_size=-1,
                  encoder_hidden_size=5,
                  max_steps=1,
                  stat_exog_list=["airline1"],
                  futr_exog_list=["trend"],
                  hist_exog_list=["y_[lag12]"],
              ),
              AutoMLP(h=12, config=config, cpus=1, num_samples=1),
              MLP(h=12, input_size=12, max_steps=1, scaler_type="robust"),
              NBEATSx(
                  h=12,
                  input_size=12,
                  max_steps=1,
                  stat_exog_list=["airline1"],
                  futr_exog_list=["trend"],
                  hist_exog_list=["y_[lag12]"],
              ),
              NHITS(h=12, input_size=12, max_steps=1, scaler_type="robust"),
              NHITS(h=12, input_size=12, loss=MQLoss(level=[80]), max_steps=1),
              TFT(h=12, input_size=24, max_steps=1, scaler_type="robust"),
              DLinear(h=12, input_size=24, max_steps=1),
              VanillaTransformer(h=12, input_size=12, max_steps=1, scaler_type=None),
              Informer(h=12, input_size=12, max_steps=1, scaler_type=None),
              Autoformer(h=12, input_size=12, max_steps=1, scaler_type=None),
              FEDformer(h=12, input_size=12, max_steps=1, scaler_type=None),
              PatchTST(h=12, input_size=24, max_steps=1, scaler_type=None),
              TimesNet(h=12, input_size=24, max_steps=1, scaler_type="standard"),
              StemGNN(h=12, input_size=12, n_series=2, max_steps=1, scaler_type="robust"),
              TSMixer(h=12, input_size=12, n_series=2, max_steps=1, scaler_type="robust"),
              TSMixerx(
                  h=12, input_size=12, n_series=2, max_steps=1, scaler_type="robust"
              ),
              DeepAR(
                  h=12,
                  input_size=24,
                  max_steps=1,
                  stat_exog_list=["airline1"],
                  futr_exog_list=["trend"],
              ),
          ],
          freq="M",
      )
      fcst.fit(df=Y_train_df, static_df=static_df)
      Y_hat_df = fcst.predict(futr_df=Y_test_df)
      Y_hat_df = Y_hat_df.merge(Y_test_df, how="left", on=["unique_id", "ds"])
      last_dates = Y_train_df.groupby("unique_id").tail(1)
      last_dates = last_dates[["unique_id", "ds"]].rename(columns={"ds": "cutoff"})
      Y_hat_df = Y_hat_df.merge(last_dates, how="left", on="unique_id")
  
      # cross validation
      fcst = NeuralForecast(
          models=[
              AutoDilatedRNN(h=12, config=config_drnn, cpus=1, num_samples=1),
              DilatedRNN(h=12, input_size=-1, encoder_hidden_size=5, max_steps=1),
              RNN(
                  h=12,
                  input_size=-1,
                  encoder_hidden_size=5,
                  max_steps=1,
                  stat_exog_list=["airline1"],
                  futr_exog_list=["trend"],
                  hist_exog_list=["y_[lag12]"],
              ),
              TCN(
                  h=12,
                  input_size=-1,
                  encoder_hidden_size=5,
                  max_steps=1,
                  stat_exog_list=["airline1"],
                  futr_exog_list=["trend"],
                  hist_exog_list=["y_[lag12]"],
              ),
              AutoMLP(h=12, config=config, cpus=1, num_samples=1),
              MLP(h=12, input_size=12, max_steps=1, scaler_type="robust"),
              NBEATSx(
                  h=12,
                  input_size=12,
                  max_steps=1,
                  stat_exog_list=["airline1"],
                  futr_exog_list=["trend"],
                  hist_exog_list=["y_[lag12]"],
              ),
              NHITS(h=12, input_size=12, max_steps=1, scaler_type="robust"),
              NHITS(h=12, input_size=12, loss=MQLoss(level=[80]), max_steps=1),
              TFT(h=12, input_size=24, max_steps=1, scaler_type="robust"),
              DLinear(h=12, input_size=24, max_steps=1),
              VanillaTransformer(h=12, input_size=12, max_steps=1, scaler_type=None),
              Informer(h=12, input_size=12, max_steps=1, scaler_type=None),
              Autoformer(h=12, input_size=12, max_steps=1, scaler_type=None),
              FEDformer(h=12, input_size=12, max_steps=1, scaler_type=None),
  