content: __all__ = ['NeuralForecast']
  
  
  import pickle
  import warnings
  from copy import deepcopy
  from itertools import chain
  from typing import Any, Dict, List, Optional, Sequence, Union
  
  import fsspec
  import numpy as np
  import pandas as pd
  import pytorch_lightning as pl
  import torch
  import utilsforecast.processing as ufp
  from coreforecast.grouped_array import GroupedArray
  from coreforecast.scalers import (
      LocalBoxCoxScaler,
      LocalMinMaxScaler,
      LocalRobustScaler,
      LocalStandardScaler,
  )
  from utilsforecast.compat import DataFrame, DFType, Series, pl_DataFrame, pl_Series
  from utilsforecast.validation import validate_freq
  from neuralforecast.common.enums import ExplainerEnum
  from neuralforecast.models import (
      GRU,
      KAN,
      LSTM,
      MLP,
      NBEATS,
      NHITS,
      RNN,
      SOFTS,
      TCN,
      TFT,
      Autoformer,
      BiTCN,
      DeepAR,
      DeepNPTS,
      DilatedRNN,
      DLinear,
      FEDformer,
      Informer,
      MLPMultivariate,
      NBEATSx,
      NLinear,
      PatchTST,
      RMoK,
      StemGNN,
      TiDE,
      TimeLLM,
      TimeMixer,
      TimesNet,
      TimeXer,
      TSMixer,
      TSMixerx,
      VanillaTransformer,
      iTransformer,
      xLSTM,
  )
  from neuralforecast.tsdataset import (
      LocalFilesTimeSeriesDataset,
      TimeSeriesDataset,
      _FilesDataset,
  )
  from neuralforecast.utils import (
      PredictionIntervals,
      get_prediction_interval_method,
      level_to_quantiles,
      quantiles_to_level,
  )
  
  from .common._base_auto import BaseAuto, MockTrial
  from .common._base_model import DistributedConfig
  from .compat import SparkDataFrame
  from .losses.pytorch import HuberIQLoss, IQLoss
  
  # this disables warnings about the number of workers in the dataloaders
  # which the user can't control
  warnings.filterwarnings("ignore", category=pl.utilities.warnings.PossibleUserWarning)
  
  
  def _insample_times(
      times: np.ndarray,
      uids: Series,
      indptr: np.ndarray,
      h: int,
      freq: Union[int, str, pd.offsets.BaseOffset],
      step_size: int = 1,
      id_col: str = "unique_id",
      time_col: str = "ds",
  ) -> DataFrame:
      sizes = np.diff(indptr)
      if (sizes < h).any():
          raise ValueError("`sizes` should be greater or equal to `h`.")
      # TODO: we can just truncate here instead of raising an error
      ns, resids = np.divmod(sizes - h, step_size)
      if (resids != 0).any():
          raise ValueError("`sizes - h` should be multiples of `step_size`")
      windows_per_serie = ns + 1
      # determine the offsets for the cutoffs, e.g. 2 means the 3rd training date is a cutoff
      cutoffs_offsets = step_size * np.hstack([np.arange(w) for w in windows_per_serie])
      # start index of each serie, e.g. [0, 17] means the the second serie starts on the 18th entry
      # we repeat each of these as many times as we have windows, e.g. windows_per_serie = [2, 3]
      # would yield [0, 0, 17, 17, 17]
      start_idxs = np.repeat(indptr[:-1], windows_per_serie)
      # determine the actual indices of the cutoffs, we repeat the cutoff for the complete horizon
      # e.g. if we have two series and h=2 this could be [0, 0, 1, 1, 17, 17, 18, 18]
      # which would have the first two training dates from each serie as the cutoffs
      cutoff_idxs = np.repeat(start_idxs + cutoffs_offsets, h)
      cutoffs = times[cutoff_idxs]
      total_windows = windows_per_serie.sum()
      # determine the offsets for the actual dates. this is going to be [0, ..., h] repeated
      ds_offsets = np.tile(np.arange(h), total_windows)
      # determine the actual indices of the times
      # e.g. if we have two series and h=2 this could be [0, 1, 1, 2, 17, 18, 18, 19]
      ds_idxs = cutoff_idxs + ds_offsets
      ds = times[ds_idxs]
      if isinstance(uids, pl_Series):
          df_constructor = pl_DataFrame
      else:
          df_constructor = pd.DataFrame
      out = df_constructor(
          {
              id_col: ufp.repeat(uids, h * windows_per_serie),
              time_col: ds,
              "cutoff": cutoffs,
          }
      )
      # the first cutoff is before the first train date
      actual_cutoffs = ufp.offset_times(out["cutoff"], freq, -1)
      out = ufp.assign_columns(out, "cutoff", actual_cutoffs)
      return out
  
  
  MODEL_FILENAME_DICT = {
      "autoformer": Autoformer,
      "autoautoformer": Autoformer,
      "deepar": DeepAR,
      "autodeepar": DeepAR,
      "dlinear": DLinear,
      "autodlinear": DLinear,
      "nlinear": NLinear,
      "autonlinear": NLinear,
      "dilatedrnn": DilatedRNN,
      "autodilatedrnn": DilatedRNN,
      "fedformer": FEDformer,
      "autofedformer": FEDformer,
      "gru": GRU,
      "autogru": GRU,
      "informer": Informer,
      "autoinformer": Informer,
      "lstm": LSTM,
      "autolstm": LSTM,
      "mlp": MLP,
      "automlp": MLP,
      "nbeats": NBEATS,
      "autonbeats": NBEATS,
      "nbeatsx": NBEATSx,
      "autonbeatsx": NBEATSx,
      "nhits": NHITS,
      "autonhits": NHITS,
      "patchtst": PatchTST,
      "autopatchtst": PatchTST,
      "rnn": RNN,
      "autornn": RNN,
      "stemgnn": StemGNN,
      "autostemgnn": StemGNN,
      "tcn": TCN,
      "autotcn": TCN,
      "tft": TFT,
      "autotft": TFT,
      "timesnet": TimesNet,
      "autotimesnet": TimesNet,
      "vanillatransformer": VanillaTransformer,
      "autovanillatransformer": VanillaTransformer,
      "timellm": TimeLLM,
      "tsmixer": TSMixer,
      "autotsmixer": TSMixer,
      "tsmixerx": TSMixerx,
      "autotsmixerx": TSMixerx,
      "mlpmultivariate": MLPMultivariate,
      "automlpmultivariate": MLPMultivariate,
      "itransformer": iTransformer,
      "autoitransformer": iTransformer,
      "bitcn": BiTCN,
      "autobitcn": BiTCN,
      "tide": TiDE,
      "autotide": TiDE,
      "deepnpts": DeepNPTS,
      "autodeepnpts": DeepNPTS,
      "softs": SOFTS,
      "autosofts": SOFTS,
      "timemixer": TimeMixer,
      "autotimemixer": TimeMixer,
      "kan": KAN,
      "autokan": KAN,
      "rmok": RMoK,
      "autormok": RMoK,
      "timexer": TimeXer,
      "autotimexer": TimeXer,
      "xlstm": xLSTM,
      "autoxlstm": xLSTM,
  }
  
  
  _type2scaler = {
      "standard": LocalStandardScaler,
      "robust": lambda: LocalRobustScaler(scale="mad"),
      "robust-iqr": lambda: LocalRobustScaler(scale="iqr"),
      "minmax": LocalMinMaxScaler,
      "boxcox": lambda: LocalBoxCoxScaler(method="loglik", lower=0.0),
  }
  
  
  class NeuralForecast:
      models: List[Any]
  
      def __init__(
          self,
          models: List[Any],
          freq: Union[str, int],
          local_scaler_type: Optional[str] = None,
          local_static_scaler_type: Optional[str] = None,
      ):
          """The `core.StatsForecast` class allows you to efficiently fit multiple `NeuralForecast` models
          for large sets of time series. It operates with pandas DataFrame `df` that identifies series
          and datestamps with the `unique_id` and `ds` columns. The `y` column denotes the target
          time series variable.
  
          Args:
              models (List[typing.Any]): Instantiated `neuralforecast.models`
                  see [collection here](./models).
              freq (str or int): Frequency of the data. Must be a valid pandas or polars offset alias, or an integer.
              local_scaler_type (str, optional): Scaler to apply per-serie to temporal features before fitting, which is inverted after predicting.
                  Can be 'standard', 'robust', 'robust-iqr', 'minmax' or 'boxcox'. Defaults to None.
              local_static_scaler_type (str, optional): Scaler to apply to static exogenous features before fitting.
                  Can be 'standard', 'robust', 'robust-iqr', 'minmax' or 'boxcox'. Defaults to None.
  
          Returns:
              NeuralForecast: Returns instantiated `NeuralForecast` class.
          """
          assert all(
              model.h == models[0].h for model in models
          ), "All models should have the same horizon"
  
          self.h = models[0].h
          self.models_init = models
          self.freq = freq
          if local_scaler_type is not None and local_scaler_type not in _type2scaler:
              raise ValueError(f"scaler_type must be one of {_type2scaler.keys()}")
          if local_static_scaler_type is not None and local_static_scaler_type not in _type2scaler:
              raise ValueError(f"static_scaler_type must be one of {_type2scaler.keys()}")
          self.local_scaler_type = local_scaler_type
          self.local_static_scaler_type = local_static_scaler_type
          self.scalers_: Dict
          self.static_scalers_: Dict
  
          # Flags and attributes
          self._fitted = False
          self._reset_models()
          self._add_level = False
  
      def _scalers_fit_transform(self, dataset: TimeSeriesDataset) -> None:
          self.scalers_, self.static_scalers_ = {}, {}
          if self.local_scaler_type is not None:
              for i, col in enumerate(dataset.temporal_cols):
                  if col == "available_mask":
                      continue
                  ga = GroupedArray(dataset.temporal[:, i].numpy(), dataset.indptr)
                  self.scalers_[col] = _type2scaler[self.local_scaler_type]().fit(ga)
                  dataset.temporal[:, i] = torch.from_numpy(self.scalers_[col].transform(ga))
          if self.local_static_scaler_type is not None and dataset.static is not None:
              for i, col in enumerate(dataset.static_cols):
                  ga = GroupedArray(dataset.static[:, i].numpy(), np.array([0, dataset.static.shape[0]]))
                  self.static_scalers_[col] = _type2scaler[self.local_static_scaler_type]().fit(ga)
                  dataset.static[:, i] = torch.from_numpy(self.static_scalers_[col].transform(ga))
  
      def _scalers_transform(self, dataset: TimeSeriesDataset) -> None:
          if self.scalers_:
              for i, col in enumerate(dataset.temporal_cols):
                  scaler = self.scalers_.get(col, None)
                  if scaler is None:
                      continue
                  ga = GroupedArray(dataset.temporal[:, i].numpy(), dataset.indptr)
                  dataset.temporal[:, i] = torch.from_numpy(scaler.transform(ga))
          if self.static_scalers_ and dataset.static is not None:
              for i, col in enumerate(dataset.static_cols):
                  scaler = self.static_scalers_.get(col, None)
                  if scaler is None:
                      continue
                  ga = GroupedArray(dataset.static[:, i].numpy(), np.array([0, dataset.static.shape[0]]))
                  dataset.static[:, i] = torch.from_numpy(scaler.transform(ga))
  
      def _scalers_target_inverse_transform(
          self, data: np.ndarray, indptr: np.ndarray
      ) -> np.ndarray:
          if not self.scalers_:
              return data
          for i in range(data.shape[1]):
              ga = GroupedArray(data[:, i], indptr)
              data[:, i] = self.scalers_[self.target_col].inverse_transform(ga)
          return data
  
      def _prepare_fit(self, df, static_df, predict_only, id_col, time_col, target_col):
          # TODO: uids, last_dates and ds should be properties of the dataset class. See github issue.
          self.id_col = id_col
          self.time_col = time_col
          self.target_col = target_col
          self._check_nan(df, static_df, id_col, time_col, target_col)
  
          dataset, uids, last_dates, ds = TimeSeriesDataset.from_df(
              df=df,
              static_df=static_df,
              id_col=id_col,
              time_col=time_col,
              target_col=target_col,
          )
          if predict_only:
              self._scalers_transform(dataset)
          else:
              self._scalers_fit_transform(dataset)
          return dataset, uids, last_dates, ds
  
      def _check_nan(self, df, static_df, id_col, time_col, target_col):
          cols_with_nans = []
  
          temporal_cols = [target_col] + [
              c for c in df.columns if c not in (id_col, time_col, target_col)
          ]
          if "available_mask" in temporal_cols:
              available_mask = df["available_mask"].to_numpy().astype(bool)
          else:
              available_mask = np.full(df.shape[0], True)
  
          df_to_check = ufp.filter_with_mask(df, available_mask)
          for col in temporal_cols:
              if ufp.is_nan_or_none(df_to_check[col]).any():
                  cols_with_nans.append(col)
  
          if static_df is not None:
              for col in [x for x in static_df.columns if x != id_col]:
                  if ufp.is_nan_or_none(static_df[col]).any():
                      cols_with_nans.append(col)
  
          if cols_with_nans:
              raise ValueError(f"Found missing values in {cols_with_nans}.")
  
      def _prepare_fit_distributed(
          self,
          df: SparkDataFrame,
          static_df: Optional[SparkDataFrame],
          id_col: str,
          time_col: str,
          target_col: str,
          distributed_config: Optional[DistributedConfig],
      ):
          if distributed_config is None:
              raise ValueError(
                  "Must set `distributed_config` when using a spark dataframe"
              )
          if self.local_scaler_type is not None:
              raise ValueError(
                  "Historic scaling isn't supported in distributed. "
                  "Please open an issue if this would be valuable to you."
              )
          if self.local_static_scaler_type is not None:
              raise ValueError(
                  "Static scaling isn't supported in distributed. "
                  "Please open an issue if this would be valuable to you."
              )
          temporal_cols = [c for c in df.columns if c not in (id_col, time_col)]
          if static_df is not None:
              static_cols = [c for c in static_df.columns if c != id_col]
              df = df.join(static_df, on=[id_col], how="left")
          else:
              static_cols = None
          self.id_col = id_col
          self.time_col = time_col
          self.target_col = target_col
          self.scalers_, self.static_scalers_ = {}, {}
          num_partitions = distributed_config.num_nodes * distributed_config.devices
          df = df.repartitionByRange(num_partitions, id_col)
          df.write.parquet(path=distributed_config.partitions_path, mode="overwrite")
          fs, _, _ = fsspec.get_fs_token_paths(distributed_config.partitions_path)
          protocol = fs.protocol
          if isinstance(protocol, tuple):
              protocol = protocol[0]
          files = [
              f"{protocol}://{file}"
              for file in fs.ls(distributed_config.partitions_path)
              if file.endswith("parquet")
          ]
          return _FilesDataset(
              files=files,
              temporal_cols=temporal_cols,
              static_cols=static_cols,
              id_col=id_col,
              time_col=time_col,
              target_col=target_col,
              min_size=df.groupBy(id_col).count().agg({"count": "min"}).first()[0],
          )
  
      def _prepare_fit_for_local_files(
          self,
          files_list: Sequence[str],
          static_df: Optional[DataFrame],
          id_col: str,
          time_col: str,
          target_col: str,
      ):
          if self.local_scaler_type is not None:
              raise ValueError(
                  "Historic scaling isn't supported when the dataset is split between files. "
                  "Please open an issue if this would be valuable to you."
              )
          if self.local_static_scaler_type is not None:
              raise ValueError(
                  "Static scaling isn't supported when the dataset is split between files. "
                  "Please open an issue if this would be valuable to you."
              )
  
          self.id_col = id_col
          self.time_col = time_col
          self.target_col = target_col
          self.scalers_, self.static_scalers_ = {}, {}
  
          exogs = self._get_needed_exog()
          return LocalFilesTimeSeriesDataset.from_data_directories(
              directories=files_list,
              static_df=static_df,
              exogs=exogs,
              id_col=id_col,
              time_col=time_col,
              target_col=target_col,
          )
  
      def fit(
          self,
          df: Optional[Union[DataFrame, SparkDataFrame, Sequence[str]]] = None,
          static_df: Optional[Union[DataFrame, SparkDataFrame]] = None,
          val_size: Optional[int] = 0,
          use_init_models: bool = False,
          verbose: bool = False,
          id_col: str = "unique_id",
          time_col: str = "ds",
          target_col: str = "y",
          distributed_config: Optional[DistributedConfig] = None,
          prediction_intervals: Optional[PredictionIntervals] = None,
      ) -> None:
          """Fit the core.NeuralForecast.
  
          Fit `models` to a large set of time series from DataFrame `df`.
          and store fitted models for later inspection.
  
          Args:
              df (pandas, polars or spark DataFrame, or a list of parquet files containing the series, optional): DataFrame with columns [`unique_id`, `ds`, `y`] and exogenous variables.
                  If None, a previously stored dataset is required. Defaults to None.
              static_df (pandas, polars or spark DataFrame, optional): DataFrame with columns [`unique_id`] and static exogenous. Defaults to None.
              val_size (int, optional): Size of validation set. Defaults to 0.
              use_init_models (bool, optional): Use initial model passed when NeuralForecast object was instantiated. Defaults to False.
              verbose (bool): Print processing steps. Defaults to False.
              id_col (str): Column that identifies each serie. Defaults to 'unique_id'.
              time_col (str): Column that identifies each timestep, its values can be timestamps or integers. Defaults to 'ds'.
              target_col (str): Column that contains the target. Defaults to 'y'.
              distributed_config (neuralforecast.DistributedConfig): Configuration to use for DDP training. Currently only spark is supported.
              prediction_intervals (PredictionIntervals, optional): Configuration to calibrate prediction intervals (Conformal Prediction). Defaults to None.
  
          Returns:
              NeuralForecast: Returns `NeuralForecast` class with fitted `models`.
          """
          if (df is None) and not (hasattr(self, "dataset")):
              raise Exception("You must pass a DataFrame or have one stored.")
  
          # Model and datasets interactions protections
          if (
              any(model.early_stop_patience_steps > 0 for model in self.models)
              and val_size == 0
          ):
              raise Exception("Set val_size>0 if early stopping is enabled.")
  
          if (val_size is not None) and (0 < val_size < self.h):
              raise ValueError(
                  f"val_size must be either 0 or greater than or equal to the horizon: {self.h}"
              )
  
          self._cs_df: Optional[DataFrame] = None
          self.prediction_intervals: Optional[PredictionIntervals] = None
  
          # Process and save new dataset (in self)
          if isinstance(df, (pd.DataFrame, pl_DataFrame)):
              validate_freq(df[time_col], self.freq)
              self.dataset, self.uids, self.last_dates, self.ds = self._prepare_fit(
                  df=df,
                  static_df=static_df,
                  predict_only=False,
                  id_col=id_col,
                  time_col=time_col,
                  target_col=target_col,
              )
              if prediction_intervals is not None:
                  self.prediction_intervals = prediction_intervals
                  self._cs_df = self._conformity_scores(
                      df=df,
                      id_col=id_col,
                      time_col=time_col,
                      target_col=target_col,
                      static_df=static_df,
                  )
  
          elif isinstance(df, SparkDataFrame):
              if static_df is not None and not isinstance(static_df, SparkDataFrame):
                  raise ValueError(
                      "`static_df` must be a spark dataframe when `df` is a spark dataframe."
                  )
              self.dataset = self._prepare_fit_distributed(
                  df=df,
                  static_df=static_df,
                  id_col=id_col,
                  time_col=time_col,
                  target_col=target_col,
                  distributed_config=distributed_config,
              )
  
              if prediction_intervals is not None:
                  raise NotImplementedError(
                      "Prediction intervals are not supported for distributed training."
                  )
  
          elif isinstance(df, Sequence):
              if not all(isinstance(val, str) for val in df):
                  raise ValueError(
                      "All entries in the list of files must be of type string"
                  )
              self.dataset = self._prepare_fit_for_local_files(
                  files_list=df,
                  static_df=static_df,
                  id_col=id_col,
                  time_col=time_col,
                  target_col=target_col,
              )
              self.uids = self.dataset.indices
              self.last_dates = self.dataset.last_times
  
              if prediction_intervals is not None:
                  raise NotImplementedError(
                      "Prediction intervals are not supported for local files."
                  )
  
          elif df is None:
              if verbose:
                  print("Using stored dataset.")
          else:
              raise ValueError(
                  f"`df` must be a pandas, polars or spark DataFrame, or a list of parquet files containing the series, or `None`, got: {type(df)}"
              )
  
          if val_size is not None:
              if self.dataset.min_size < val_size:
                  warnings.warn(
                      "Validation set size is larger than the shorter time-series."
                  )
  
          # Recover initial model if use_init_models
          if use_init_models:
              self._reset_models()
  
          for i, model in enumerate(self.models):
              self.models[i] = model.fit(
                  self.dataset, val_size=val_size, distributed_config=distributed_config
              )
  
          self._fitted = True
  
      def make_future_dataframe(
          self, df: Optional[DFType] = None, h: Optional[int] = None
      ) -> DFType:
          """Create a dataframe with all ids and future times in the forecasting horizon.
  
          Args:
              df (pandas or polars DataFrame, optional): DataFrame with columns [`unique_id`, `ds`, `y`] and exogenous variables.
                  Only required if this is different than the one used in the fit step. Defaults to None.
          """
          if not self._fitted:
              raise Exception("You must fit the model first.")
          if df is not None:
              df = ufp.sort(df, by=[self.id_col, self.time_col])
              last_times_by_id = ufp.group_by_agg(
                  df,
                  by=self.id_col,
                  aggs={self.time_col: "max"},
                  maintain_order=True,
              )
              uids = last_times_by_id[self.id_col]
              last_times = last_times_by_id[self.time_col]
          else:
              uids = self.uids
              last_times = self.last_dates
          if h is None:
              h = self.h
          return ufp.make_future_dataframe(
              uids=uids,
              last_times=last_times,
              freq=self.freq,
              h=h,
              id_col=self.id_col,
              time_col=self.time_col,
          )
  
      def get_missing_future(
          self, futr_df: DFType, df: Optional[DFType] = None, h: Optional[int] = None
      ) -> DFType:
          """Get the missing ids and times combinations in `futr_df`.
  
          Args:
              futr_df (pandas or polars DataFrame): DataFrame with [`unique_id`, `ds`] columns and `df`'s future exogenous.
              df (pandas or polars DataFrame, optional): DataFrame with columns [`unique_id`, `ds`, `y`] and exogenous variables.
                  Only required if this is different than the one used in the fit step. Defaults to None.
          """
          expected = self.make_future_dataframe(df, h=h)
          ids = [self.id_col, self.time_col]
          return ufp.anti_join(expected, futr_df[ids], on=ids)
  
      def _get_needed_futr_exog(self):
          futr_exogs = []
          for m in self.models:
              if isinstance(m, BaseAuto):
                  if isinstance(m.config, dict):  # ray
                      exogs = m.config.get("futr_exog_list", [])
                      if hasattr(
                          exogs, "categories"
                      ):  # features are being tuned, get possible values
                          exogs = exogs.categories
                  else:  # optuna
                      exogs = m.config(MockTrial()).get("futr_exog_list", [])
              else:  # regular model, extract them directly
                  exogs = getattr(m, "futr_exog_list", [])
  
              for exog in exogs:
                  if isinstance(exog, str):
                      futr_exogs.append(exog)
                  else:
                      futr_exogs.extend(exog)
  
          return set(futr_exogs)
  
      def _get_needed_exog(self):
          futr_exog = self._get_needed_futr_exog()
  
          hist_exog = []
          for m in self.models:
              if isinstance(m, BaseAuto):
                  if isinstance(m.config, dict):  # ray
                      exogs = m.config.get("hist_exog_list", [])
                      if hasattr(
                          exogs, "categories"
                      ):  # features are being tuned, get possible values
                          exogs = exogs.categories
                  else:  # optuna
                      exogs = m.config(MockTrial()).get("hist_exog_list", [])
              else:  # regular model, extract them directly
                  exogs = getattr(m, "hist_exog_list", [])
  
              for exog in exogs:
                  if isinstance(exog, str):
                      hist_exog.append(exog)
                  else:
                      hist_exog.extend(exog)
  
          return futr_exog | set(hist_exog)
  
      def _get_model_names(self, add_level=False) -> List[str]:
          names: List[str] = []
          count_names = {"model": 0}
          for model in self.models:
              model_name = repr(model)
              count_names[model_name] = count_names.get(model_name, -1) + 1
              if count_names[model_name] > 0:
                  model_name += str(count_names[model_name])
  
              if add_level and (
                  model.loss.outputsize_multiplier > 1
                  or isinstance(model.loss, (IQLoss, HuberIQLoss))
              ):
                  continue
  
              names.extend(model_name + n for n in model.loss.output_names)
          return names
  
      def _predict_distributed(
          self,
          df: Optional[SparkDataFrame],
          static_df: Optional[SparkDataFrame],
          futr_df: Optional[SparkDataFrame],
          engine,
          h: Optional[int] = None,
      ):
          import fugue.api as fa
  
          def _predict(
              df: pd.DataFrame,
              static_cols,
              futr_exog_cols,
              models,
              freq,
              id_col,
              time_col,
              target_col,
              h,
          ) -> pd.DataFrame:
              from neuralforecast import NeuralForecast
  
              nf = NeuralForecast(models=models, freq=freq)
              nf.id_col = id_col
              nf.time_col = time_col
              nf.target_col = target_col
              nf.scalers_ = {}
              nf._fitted = True
              if futr_exog_cols:
                  # if we have futr_exog we'll have extra rows with the future values
                  futr_rows = df[target_col].isnull()
                  futr_df = df.loc[
                      futr_rows, [self.id_col, self.time_col] + futr_exog_cols
                  ].copy()
                  df = df[~futr_rows].copy()
              else:
                  futr_df = None
              if static_cols:
                  static_df = (
                      df[[self.id_col] + static_cols]
                      .groupby(self.id_col, observed=True)
                      .head(1)
                  )
                  df = df.drop(columns=static_cols)
              else:
                  static_df = None
              return nf.predict(df=df, static_df=static_df, futr_df=futr_df, h=h)
  
          # df
          if isinstance(df, SparkDataFrame):
              repartition = True
          else:
              if engine is None:
                  raise ValueError("engine is required for distributed inference")
              df = engine.read.parquet(*self.dataset.files)
              # we save the datataset with partitioning
              repartition = False
  
          # static
          static_cols = set(
              chain.from_iterable(getattr(m, "stat_exog_list", []) for m in self.models)
          )
          if static_df is not None:
              if not isinstance(static_df, SparkDataFrame):
                  raise ValueError(
                      "`static_df` must be a spark dataframe when `df` is a spark dataframe "
                      "or the models were trained in a distributed setting.\n"
                      "You can also provide local dataframes (pandas or polars) as `df` and `static_df`."
                  )
              missing_static = static_cols - set(static_df.columns)
              if missing_static:
                  raise ValueError(
                      f"The following static columns are missing from the static_df: {missing_static}"
                  )
              # join is supposed to preserve the partitioning
              df = df.join(static_df, on=[self.id_col], how="left")
  
          # exog
          if futr_df is not None:
              if not isinstance(futr_df, SparkDataFrame):
                  raise ValueError(
                      "`futr_df` must be a spark dataframe when `df` is a spark dataframe "
                      "or the models were trained in a distributed setting.\n"
                      "You can also provide local dataframes (pandas or polars) as `df` and `futr_df`."
                  )
              if self.target_col in futr_df.columns:
                  raise ValueError("`futr_df` must not contain the target column.")
              # df has the statics, historic exog and target at this point, futr_df doesnt
              df = df.unionByName(futr_df, allowMissingColumns=True)
              # union doesn't guarantee preserving the partitioning
              repartition = True
  
          if repartition:
              df = df.repartitionByRange(df.rdd.getNumPartitions(), self.id_col)
  
          # predict
          base_schema = fa.get_schema(df).extract([self.id_col, self.time_col])
          models_schema = {model: "float" for model in self._get_model_names()}
          return fa.transform(
              df=df,
              using=_predict,
              schema=base_schema.append(models_schema),
              params=dict(
                  static_cols=list(static_cols),
                  futr_exog_cols=list(self._get_needed_futr_exog()),
                  models=self.models,
                  freq=self.freq,
                  id_col=self.id_col,
                  time_col=self.time_col,
                  target_col=self.target_col,
                  h=h,
              ),
          )
  
      def predict(
          self,
          df: Optional[Union[DataFrame, SparkDataFrame]] = None,
          static_df: Optional[Union[DataFrame, SparkDataFrame]] = None,
          futr_df: Optional[Union[DataFrame, SparkDataFrame]] = None,
          verbose: bool = False,
          engine=None,
          level: Optional[List[Union[int, float]]] = None,
          quantiles: Optional[List[float]] = None,
          h: Optional[int] = None,
          **data_kwargs,
      ):
          """Predict with core.NeuralForecast.
  
          Use stored fitted `models` to predict large set of time series from DataFrame `df`.
  
          Args:
              df (pandas, polars or spark DataFrame, optional): DataFrame with columns [`unique_id`, `ds`, `y`] and exogenous variables.
                  If a DataFrame is passed, it is used to generate forecasts. Defaults to None.
              static_df (pandas, polars or spark DataFrame, optional): DataFrame with columns [`unique_id`] and static exogenous. Defaults to None.
              futr_df (pandas, polars or spark DataFrame, optional): DataFrame with [`unique_id`, `ds`] columns and `df`'s future exogenous. Defaults to None.
              verbose (bool): Print processing steps. Defaults to False.
              engine (spark session): Distributed engine for inference. Only used if df is a spark dataframe or if fit was called on a spark dataframe.
              level (list of ints or floats, optional): Confidence levels between 0 and 100. Defaults to None.
              quantiles (list of floats, optional): Alternative to level, target quantiles to predict. Defaults to None.
              h (int, optional): Forecasting horizon. If None, uses the horizon of the fitted models. Defaults to None.
              data_kwargs (kwargs): Extra arguments to be passed to the dataset within each model.
  
          Returns:
              fcsts_df (pandas or polars DataFrame): DataFrame with insample `models` columns for point predictions and probabilistic
                  predictions for all fitted `models`.
  