content: __all__ = ['BaseAuto']
  
  
  import warnings
  from copy import deepcopy
  from os import cpu_count
  
  import pytorch_lightning as pl
  import torch
  from ray import air, tune
  from ray.tune.integration.pytorch_lightning import TuneReportCallback
  from ray.tune.search.basic_variant import BasicVariantGenerator
  
  
  class MockTrial:
      def suggest_int(*args, **kwargs):
          return "int"
  
      def suggest_categorical(self, name, choices):
          return choices
  
      def suggest_uniform(*args, **kwargs):
          return "uniform"
  
      def suggest_loguniform(*args, **kwargs):
          return "loguniform"
  
      def suggest_float(*args, **kwargs):
          if "log" in kwargs:
              return "quantized_log"
          elif "step" in kwargs:
              return "quantized_loguniform"
          return "float"
  
  
  class BaseAuto(pl.LightningModule):
      """
      Class for Automatic Hyperparameter Optimization, it builds on top of `ray` to
      give access to a wide variety of hyperparameter optimization tools ranging
      from classic grid search, to Bayesian optimization and HyperBand algorithm.
  
      The validation loss to be optimized is defined by the `config['loss']` dictionary
      value, the config also contains the rest of the hyperparameter search space.
  
      It is important to note that the success of this hyperparameter optimization
      heavily relies on a strong correlation between the validation and test periods.
  
      Args:
          cls_model (PyTorch/PyTorchLightning model): See `neuralforecast.models` [collection here](./models).
          h (int): Forecast horizon
          loss (PyTorch module): Instantiated train loss class from [losses collection](./losses.pytorch).
          valid_loss (PyTorch module): Instantiated valid loss class from [losses collection](./losses.pytorch).
          config (dict or callable): Dictionary with ray.tune defined search space or function that takes an optuna trial and returns a configuration dict.
          search_alg (ray.tune.search variant or optuna.sampler): For ray see https://docs.ray.io/en/latest/tune/api_docs/suggestion.html
              For optuna see https://optuna.readthedocs.io/en/stable/reference/samplers/index.html.
          num_samples (int): Number of hyperparameter optimization steps/samples.
          cpus (int): Number of cpus to use during optimization. Only used with ray tune.
          gpus (int): Number of gpus to use during optimization, default all available. Only used with ray tune.
          refit_with_val (bool): Refit of best model should preserve val_size.
          verbose (bool): Track progress.
          alias (str): Custom name of the model.
          backend (str): Backend to use for searching the hyperparameter space, can be either 'ray' or 'optuna'.
          callbacks (list of callable): List of functions to call during the optimization process.
              ray reference: https://docs.ray.io/en/latest/tune/tutorials/tune-metrics.html
              optuna reference: https://optuna.readthedocs.io/en/stable
      """
  
      def __init__(
          self,
          cls_model,
          h,
          loss,
          valid_loss,
          config,
          search_alg=BasicVariantGenerator(random_state=1),
          num_samples=10,
          cpus=cpu_count(),
          gpus=torch.cuda.device_count(),
          refit_with_val=False,
          verbose=False,
          alias=None,
          backend="ray",
          callbacks=None,
      ):
          super(BaseAuto, self).__init__()
          with warnings.catch_warnings(record=False):
              warnings.filterwarnings("ignore")
              # the following line issues a warning about the loss attribute being saved
              # but we do want to save it
              self.save_hyperparameters()  # Allows instantiation from a checkpoint from class
  
          if backend == "ray":
              if not isinstance(config, dict):
                  raise ValueError(
                      "You have to provide a dict as `config` when using `backend='ray'`"
                  )
              config_base = deepcopy(config)
          elif backend == "optuna":
              if not callable(config):
                  raise ValueError(
                      "You have to provide a function that takes a trial and returns a dict as `config` when using `backend='optuna'`"
                  )
              # extract constant values from the config fn for validations
              config_base = config(MockTrial())
          else:
              raise ValueError(
                  f"Unknown backend {backend}. The supported backends are 'ray' and 'optuna'."
              )
          if config_base.get("h", None) is not None:
              raise Exception("Please use `h` init argument instead of `config['h']`.")
          if config_base.get("loss", None) is not None:
              raise Exception(
                  "Please use `loss` init argument instead of `config['loss']`."
              )
          if config_base.get("valid_loss", None) is not None:
              raise Exception(
                  "Please use `valid_loss` init argument instead of `config['valid_loss']`."
              )
          # This attribute helps to protect
          # model and datasets interactions protections
          if "early_stop_patience_steps" in config_base.keys():
              self.early_stop_patience_steps = 1
          else:
              self.early_stop_patience_steps = -1
  
          if callable(config):
              # reset config_base here to save params to override in the config fn
              config_base = {}
  
          # Add losses to config and protect valid_loss default
          config_base["h"] = h
          config_base["loss"] = loss
          if valid_loss is None:
              valid_loss = loss
          config_base["valid_loss"] = valid_loss
  
          if isinstance(config, dict):
              self.config = config_base
          else:
  
              def config_f(trial):
                  return {**config(trial), **config_base}
  
              self.config = config_f
  
          self.h = h
          self.cls_model = cls_model
          self.loss = loss
          self.valid_loss = valid_loss
  
          self.num_samples = num_samples
          self.search_alg = search_alg
          self.cpus = cpus
          self.gpus = gpus
          self.refit_with_val = refit_with_val or self.early_stop_patience_steps > 0
          self.verbose = verbose
          self.alias = alias
          self.backend = backend
          self.callbacks = callbacks
  
          # Base Class attributes
          self.EXOGENOUS_FUTR = cls_model.EXOGENOUS_FUTR
          self.EXOGENOUS_HIST = cls_model.EXOGENOUS_HIST
          self.EXOGENOUS_STAT = cls_model.EXOGENOUS_STAT
          self.MULTIVARIATE = cls_model.MULTIVARIATE
          self.RECURRENT = cls_model.RECURRENT
  
      def __repr__(self):
          return type(self).__name__ if self.alias is None else self.alias
  
      def _train_tune(self, config_step, cls_model, dataset, val_size, test_size):
          """BaseAuto._train_tune
  
          Internal function that instantiates a NF class model, then automatically
          explores the validation loss (ptl/val_loss) on which the hyperparameter
          exploration is based.
  
          Args:
              config_step (dict): Dict, initialization parameters of a NF model.
              cls_model (NeuralForecast model class): NeuralForecast model class, yet to be instantiated.
              dataset (NeuralForecast dataset): NeuralForecast dataset, to fit the model.
              val_size (int): Validation size for temporal cross-validation.
              test_size (int): Test size for temporal cross-validation.
          """
          metrics = {"loss": "ptl/val_loss", "train_loss": "train_loss"}
          callbacks = [TuneReportCallback(metrics, on="validation_end")]
          if "callbacks" in config_step.keys():
              callbacks.extend(config_step["callbacks"])
          config_step = {**config_step, **{"callbacks": callbacks}}
  
          # Protect dtypes from tune samplers
          if "batch_size" in config_step.keys():
              config_step["batch_size"] = int(config_step["batch_size"])
          if "windows_batch_size" in config_step.keys():
              config_step["windows_batch_size"] = int(config_step["windows_batch_size"])
  
          # Tune session receives validation signal
          # from the specialized PL TuneReportCallback
          _ = self._fit_model(
              cls_model=cls_model,
              config=config_step,
              dataset=dataset,
              val_size=val_size,
              test_size=test_size,
          )
  
      def _tune_model(
          self,
          cls_model,
          dataset,
          val_size,
          test_size,
          cpus,
          gpus,
          verbose,
          num_samples,
          search_alg,
          config,
      ):
          train_fn_with_parameters = tune.with_parameters(
              self._train_tune,
              cls_model=cls_model,
              dataset=dataset,
              val_size=val_size,
              test_size=test_size,
          )
  
          # Device
          if gpus > 0:
              device_dict = {"gpu": gpus}
          else:
              device_dict = {"cpu": cpus}
  
          # on Windows, prevent long trial directory names
          import platform
  
          trial_dirname_creator = (
              (lambda trial: f"{trial.trainable_name}_{trial.trial_id}")
              if platform.system() == "Windows"
              else None
          )
  
          tuner = tune.Tuner(
              tune.with_resources(train_fn_with_parameters, device_dict),
              run_config=air.RunConfig(callbacks=self.callbacks, verbose=verbose),
              tune_config=tune.TuneConfig(
                  metric="loss",
                  mode="min",
                  num_samples=num_samples,
                  search_alg=search_alg,
                  trial_dirname_creator=trial_dirname_creator,
              ),
              param_space=config,
          )
          results = tuner.fit()
          return results
  
      @staticmethod
      def _ray_config_to_optuna(ray_config):
          def optuna_config(trial):
              out = {}
              for k, v in ray_config.items():
                  if hasattr(v, "sampler"):
                      sampler = v.sampler
                      if isinstance(
                          sampler, tune.search.sample.Integer.default_sampler_cls
                      ):
                          v = trial.suggest_int(k, v.lower, v.upper)
                      elif isinstance(
                          sampler, tune.search.sample.Categorical.default_sampler_cls
                      ):
                          v = trial.suggest_categorical(k, v.categories)
                      elif isinstance(sampler, tune.search.sample.Uniform):
                          v = trial.suggest_uniform(k, v.lower, v.upper)
                      elif isinstance(sampler, tune.search.sample.LogUniform):
                          v = trial.suggest_loguniform(k, v.lower, v.upper)
                      elif isinstance(sampler, tune.search.sample.Quantized):
                          if isinstance(
                              sampler.get_sampler(), tune.search.sample.Float._LogUniform
                          ):
                              v = trial.suggest_float(k, v.lower, v.upper, log=True)
                          elif isinstance(
                              sampler.get_sampler(), tune.search.sample.Float._Uniform
                          ):
                              v = trial.suggest_float(k, v.lower, v.upper, step=sampler.q)
                      else:
                          raise ValueError(f"Couldn't translate {type(v)} to optuna.")
                  out[k] = v
              return out
  
          return optuna_config
  
      def _optuna_tune_model(
          self,
          cls_model,
          dataset,
          val_size,
          test_size,
          verbose,
          num_samples,
          search_alg,
          config,
          distributed_config,
      ):
          import optuna
  
          def objective(trial):
              user_cfg = config(trial)
              cfg = deepcopy(user_cfg)
              model = self._fit_model(
                  cls_model=cls_model,
                  config=cfg,
                  dataset=dataset,
                  val_size=val_size,
                  test_size=test_size,
                  distributed_config=distributed_config,
              )
              trial.set_user_attr("ALL_PARAMS", user_cfg)
              metrics = model.metrics
              trial.set_user_attr(
                  "METRICS",
                  {
                      "loss": metrics["ptl/val_loss"],
                      "train_loss": metrics["train_loss"],
                  },
              )
              return trial.user_attrs["METRICS"]["loss"]
  
          if isinstance(search_alg, optuna.samplers.BaseSampler):
              sampler = search_alg
          else:
              sampler = None
  
          study = optuna.create_study(sampler=sampler, direction="minimize")
          study.optimize(
              objective,
              n_trials=num_samples,
              show_progress_bar=verbose,
              callbacks=self.callbacks,
          )
          return study
  
      def _fit_model(
          self, cls_model, config, dataset, val_size, test_size, distributed_config=None
      ):
          model = cls_model(**config)
          model = model.fit(
              dataset,
              val_size=val_size,
              test_size=test_size,
              distributed_config=distributed_config,
          )
          return model
  
      def fit(
          self,
          dataset,
          val_size=0,
          test_size=0,
          random_seed=None,
          distributed_config=None,
      ):
          """BaseAuto.fit
  
          Perform the hyperparameter optimization as specified by the BaseAuto configuration
          dictionary `config`.
  
          The optimization is performed on the `TimeSeriesDataset` using temporal cross validation with
          the validation set that sequentially precedes the test set.
  
          Args:
              dataset (NeuralForecast's `TimeSeriesDataset`): NeuralForecast's `TimeSeriesDataset` see details [here](./tsdataset)
              val_size (int): Size of temporal validation set (needs to be bigger than 0).
              test_size (int): Size of temporal test set (default 0).
              random_seed (int): Random seed for hyperparameter exploration algorithms, not yet implemented.
  
          Returns:
              self: Fitted instance of `BaseAuto` with best hyperparameters and results.
          """
          # we need val_size > 0 to perform
          # hyperparameter selection.
          search_alg = deepcopy(self.search_alg)
          val_size = val_size if val_size > 0 else self.h
          if self.backend == "ray":
              if distributed_config is not None:
                  raise ValueError(
                      "distributed training is not supported for the ray backend."
                  )
              results = self._tune_model(
                  cls_model=self.cls_model,
                  dataset=dataset,
                  val_size=val_size,
                  test_size=test_size,
                  cpus=self.cpus,
                  gpus=self.gpus,
                  verbose=self.verbose,
                  num_samples=self.num_samples,
                  search_alg=search_alg,
                  config=self.config,
              )
              best_config = results.get_best_result().config
          else:
              results = self._optuna_tune_model(
                  cls_model=self.cls_model,
                  dataset=dataset,
                  val_size=val_size,
                  test_size=test_size,
                  verbose=self.verbose,
                  num_samples=self.num_samples,
                  search_alg=search_alg,
                  config=self.config,
                  distributed_config=distributed_config,
              )
              best_config = results.best_trial.user_attrs["ALL_PARAMS"]
          self.model = self._fit_model(
              cls_model=self.cls_model,
              config=best_config,
              dataset=dataset,
              val_size=val_size * self.refit_with_val,
              test_size=test_size,
              distributed_config=distributed_config,
          )
          self.results = results
  
          # Added attributes for compatibility with NeuralForecast core
          self.futr_exog_list = self.model.futr_exog_list
          self.hist_exog_list = self.model.hist_exog_list
          self.stat_exog_list = self.model.stat_exog_list
          return self
  
      def predict(self, dataset, step_size=1, h=None, **data_kwargs):
          """BaseAuto.predict
  
          Predictions of the best performing model on validation.
  
          Args:
              dataset (NeuralForecast's `TimeSeriesDataset`): NeuralForecast's `TimeSeriesDataset` see details [here](./tsdataset)
              step_size (int): Steps between sequential predictions, (default 1).
              h (int): Prediction horizon, if None, uses the model's fitted horizon. Defaults to None.
              **data_kwarg: Additional parameters for the dataset module.
  
          Returns:
              y_hat: Numpy predictions of the `NeuralForecast` model.
          """
          return self.model.predict(dataset=dataset, step_size=step_size, h=h, **data_kwargs)
  
      def set_test_size(self, test_size):
          self.model.set_test_size(test_size)
  
      def get_test_size(self):
          return self.model.test_size
  
      def save(self, path):
          """BaseAuto.save
  
          Save the fitted model to disk.
  
          Args:
              path (str): Path to save the model.
          """
          self.model.save(path)
  