# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/core.ipynb.

# %% auto 0
__all__ = ['NeuralForecast']

# %% ../nbs/core.ipynb 4
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 .common._base_model import DistributedConfig
from .compat import SparkDataFrame
from neuralforecast.tsdataset import (
    _FilesDataset,
    TimeSeriesDataset,
    LocalFilesTimeSeriesDataset,
)
from neuralforecast.models import (
    GRU,
    LSTM,
    RNN,
    TCN,
    DeepAR,
    DilatedRNN,
    MLP,
    NHITS,
    NBEATS,
    NBEATSx,
    DLinear,
    NLinear,
    TFT,
    VanillaTransformer,
    Informer,
    Autoformer,
    FEDformer,
    StemGNN,
    PatchTST,
    TimesNet,
    TimeLLM,
    TSMixer,
    TSMixerx,
    MLPMultivariate,
    iTransformer,
    BiTCN,
    TiDE,
    DeepNPTS,
    SOFTS,
    TimeMixer,
    KAN,
    RMoK,
)
from .common._base_auto import BaseAuto, MockTrial
from .utils import PredictionIntervals, get_prediction_interval_method

# %% ../nbs/core.ipynb 5
# 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

# %% ../nbs/core.ipynb 7
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,
}

# %% ../nbs/core.ipynb 8
_type2scaler = {
    "standard": LocalStandardScaler,
    "robust": lambda: LocalRobustScaler(scale="mad"),
    "robust-iqr": lambda: LocalRobustScaler(scale="iqr"),
    "minmax": LocalMinMaxScaler,
    "boxcox": lambda: LocalBoxCoxScaler(method="loglik", lower=0.0),
}

# %% ../nbs/core.ipynb 9
class NeuralForecast:

    def __init__(
        self,
        models: List[Any],
        freq: Union[str, int],
        local_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.

        Parameters
        ----------
        models : List[typing.Any]
            Instantiated `neuralforecast.models`
            see [collection here](https://nixtla.github.io/neuralforecast/models.html).
        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 (default=None)
            Scaler to apply per-serie to all features before fitting, which is inverted after predicting.
            Can be 'standard', 'robust', 'robust-iqr', 'minmax' or 'boxcox'

        Returns
        -------
        self : 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()}")
        self.local_scaler_type = local_scaler_type
        self.scalers_: Dict

        # Flags and attributes
        self._fitted = False
        self._reset_models()

    def _scalers_fit_transform(self, dataset: TimeSeriesDataset) -> None:
        self.scalers_ = {}
        if self.local_scaler_type is None:
            return 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))

    def _scalers_transform(self, dataset: TimeSeriesDataset) -> None:
        if not self.scalers_:
            return None
        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))

    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."
            )
        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_ = {}
        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."
            )

        self.id_col = id_col
        self.time_col = time_col
        self.target_col = target_col
        self.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.

        Parameters
        ----------
        df : pandas, polars or spark DataFrame, or a list of parquet files containing the series, optional (default=None)
            DataFrame with columns [`unique_id`, `ds`, `y`] and exogenous variables.
            If None, a previously stored dataset is required.
        static_df : pandas, polars or spark DataFrame, optional (default=None)
            DataFrame with columns [`unique_id`] and static exogenous.
        val_size : int, optional (default=0)
            Size of validation set.
        use_init_models : bool, optional (default=False)
            Use initial model passed when NeuralForecast object was instantiated.
        verbose : bool (default=False)
            Print processing steps.
        id_col : str (default='unique_id')
            Column that identifies each serie.
        time_col : str (default='ds')
            Column that identifies each timestep, its values can be timestamps or integers.
        target_col : str (default='y')
            Column that contains the target.
        distributed_config : neuralforecast.DistributedConfig
            Configuration to use for DDP training. Currently only spark is supported.
        prediction_intervals : PredictionIntervals, optional (default=None)
            Configuration to calibrate prediction intervals (Conformal Prediction).

        Returns
        -------
        self : 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.")

        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) -> DFType:
        """Create a dataframe with all ids and future times in the forecasting horizon.

        Parameters
        ----------
        df : pandas or polars DataFrame, optional (default=None)
            DataFrame with columns [`unique_id`, `ds`, `y`] and exogenous variables.
            Only required if this is different than the one used in the fit step.
        """
        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
        return ufp.make_future_dataframe(
            uids=uids,
            last_times=last_times,
            freq=self.freq,
            h=self.h,
            id_col=self.id_col,
            time_col=self.time_col,
        )

    def get_missing_future(
        self, futr_df: DFType, df: Optional[DFType] = None
    ) -> DFType:
        """Get the missing ids and times combinations in `futr_df`.

        Parameters
        ----------
        futr_df : pandas or polars DataFrame
            DataFrame with [`unique_id`, `ds`] columns and `df`'s future exogenous.
        df : pandas or polars DataFrame, optional (default=None)
            DataFrame with columns [`unique_id`, `ds`, `y`] and exogenous variables.
            Only required if this is different than the one used in the fit step.
        """
        expected = self.make_future_dataframe(df)
        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:
            if add_level and model.loss.outputsize_multiplier > 1:
                continue

            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])
            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,
    ):
        import fugue.api as fa

        def _predict(
            df: pd.DataFrame,
            static_cols,
            futr_exog_cols,
            models,
            freq,
            id_col,
            time_col,
            target_col,
        ) -> 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)

        # 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,
            ),
        )

    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,
        **data_kwargs,
    ):
        """Predict with core.NeuralForecast.

        Use stored fitted `models` to predict large set of time series from DataFrame `df`.

        Parameters
        ----------
        df : pandas, polars or spark DataFrame, optional (default=None)
            DataFrame with columns [`unique_id`, `ds`, `y`] and exogenous variables.
            If a DataFrame is passed, it is used to generate forecasts.
        static_df : pandas, polars or spark DataFrame, optional (default=None)
            DataFrame with columns [`unique_id`] and static exogenous.
        futr_df : pandas, polars or spark DataFrame, optional (default=None)
            DataFrame with [`unique_id`, `ds`] columns and `df`'s future exogenous.
        verbose : bool (default=False)
            Print processing steps.
        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 (default=None)
            Confidence levels between 0 and 100.
        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`.
        """
        if df is None and not hasattr(self, "dataset"):
            raise Exception("You must pass a DataFrame or have one stored.")

        if not self._fitted:
            raise Exception("You must fit the model before predicting.")

        needed_futr_exog = self._get_needed_futr_exog()
        if needed_futr_exog:
            if futr_df is None:
                raise ValueError(
                    f"Models require the following future exogenous features: {needed_futr_exog}. "
                    "Please provide them through the `futr_df` argument."
                )
            else:
                missing = needed_futr_exog - set(futr_df.columns)
                if missing:
                    raise ValueError(
                        f"The following features are missing from `futr_df`: {missing}"
                    )

        # distributed df or NeuralForecast instance was trained with a distributed input and no df is provided
        # we assume the user wants to perform distributed inference as well
        is_files_dataset = isinstance(getattr(self, "dataset", None), _FilesDataset)
        is_dataset_local_files = isinstance(
            getattr(self, "dataset", None), LocalFilesTimeSeriesDataset
        )
        if isinstance(df, SparkDataFrame) or (df is None and is_files_dataset):
            return self._predict_distributed(
                df=df,
                static_df=static_df,
                futr_df=futr_df,
                engine=engine,
            )

        if is_dataset_local_files and df is None:
            raise ValueError(
                "When the model has been trained on a dataset that is split between multiple files, you must pass in a specific dataframe for prediciton."
            )

        # Process new dataset but does not store it.
        if df is not None:
            validate_freq(df[self.time_col], self.freq)
            dataset, uids, last_dates, _ = self._prepare_fit(
                df=df,
                static_df=static_df,
                predict_only=True,
                id_col=self.id_col,
                time_col=self.time_col,
                target_col=self.target_col,
            )
        else:
            dataset = self.dataset
            uids = self.uids
            last_dates = self.last_dates
            if verbose:
                print("Using stored dataset.")

        cols = self._get_model_names()

        # Placeholder dataframe for predictions with unique_id and ds
        fcsts_df = ufp.make_future_dataframe(
            uids=uids,
            last_times=last_dates,
            freq=self.freq,
            h=self.h,
            id_col=self.id_col,
            time_col=self.time_col,
        )

        # Update and define new forecasting dataset
        if futr_df is None:
            futr_df = fcsts_df
        else:
            futr_orig_rows = futr_df.shape[0]
            futr_df = ufp.join(futr_df, fcsts_df, on=[self.id_col, self.time_col])
            if futr_df.shape[0] < fcsts_df.shape[0]:
                if df is None:
                    expected_cmd = "make_future_dataframe()"
                    missing_cmd = "get_missing_future(futr_df)"
                else:
                    expected_cmd = "make_future_dataframe(df)"
                    missing_cmd = "get_missing_future(futr_df, df)"
                raise ValueError(
                    "There are missing combinations of ids and times in `futr_df`.\n"
                    f"You can run the `{expected_cmd}` method to get the expected combinations or "
                    f"the `{missing_cmd}` method to get the missing combinations."
                )
            if futr_orig_rows > futr_df.shape[0]:
                dropped_rows = futr_orig_rows - futr_df.shape[0]
                warnings.warn(f"Dropped {dropped_rows:,} unused rows from `futr_df`.")
            if any(ufp.is_none(futr_df[col]).any() for col in needed_futr_exog):
                raise ValueError("Found null values in `futr_df`")
        futr_dataset = dataset.align(
            futr_df,
            id_col=self.id_col,
            time_col=self.time_col,
            target_col=self.target_col,
        )
        self._scalers_transform(futr_dataset)
        dataset = dataset.append(futr_dataset)

        col_idx = 0
        fcsts = np.full(
            (self.h * len(uids), len(cols)), fill_value=np.nan, dtype=np.float32
        )
        for model in self.models:
            old_test_size = model.get_test_size()
            model.set_test_size(self.h)  # To predict h steps ahead
            model_fcsts = model.predict(dataset=dataset, **data_kwargs)
            # Append predictions in memory placeholder
            output_length = len(model.loss.output_names)
            fcsts[:, col_idx : col_idx + output_length] = model_fcsts
            col_idx += output_length
            model.set_test_size(old_test_size)  # Set back to original value
        if self.scalers_:
            indptr = np.append(0, np.full(len(uids), self.h).cumsum())
            fcsts = self._scalers_target_inverse_transform(fcsts, indptr)

        # Declare predictions pd.DataFrame
        cols = (
            self._get_model_names()
        )  # Needed for IQLoss as column names may have changed during the call to .predict()
        if isinstance(fcsts_df, pl_DataFrame):
            fcsts = pl_DataFrame(dict(zip(cols, fcsts.T)))
        else:
            fcsts = pd.DataFrame(fcsts, columns=cols)
        fcsts_df = ufp.horizontal_concat([fcsts_df, fcsts])

        # add prediction intervals
        if level is not None:
            if self._cs_df is None or self.prediction_intervals is None:
                raise Exception(
                    "You must fit the model with prediction_intervals to use level."
                )
            else:
                level_ = sorted(level)
                model_names = self._get_model_names(add_level=True)
                prediction_interval_method = get_prediction_interval_method(
                    self.prediction_intervals.method
                )

                fcsts_df = prediction_interval_method(
                    fcsts_df,
                    self._cs_df,
                    model_names=list(model_names),
                    level=level_,
                    cs_n_windows=self.prediction_intervals.n_windows,
                    n_series=len(uids),
                    horizon=self.h,
                )

        return fcsts_df

    def _reset_models(self):
        self.models = [deepcopy(model) for model in self.models_init]
        if self._fitted:
            print("WARNING: Deleting previously fitted models.")

    def _no_refit_cross_validation(
        self,
        df: Optional[DataFrame],
        static_df: Optional[DataFrame],
        n_windows: int,
        step_size: int,
        val_size: Optional[int],
        test_size: int,
        verbose: bool,
        id_col: str,
        time_col: str,
        target_col: str,
        **data_kwargs,
    ) -> DataFrame:
        if (df is None) and not (hasattr(self, "dataset")):
            raise Exception("You must pass a DataFrame or have one stored.")

        # Process and save new dataset (in self)
        if df is not None:
            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,
            )
        else:
            if verbose:
                print("Using stored dataset.")

        if val_size is not None:
            if self.dataset.min_size < (val_size + test_size):
                warnings.warn(
                    "Validation and test sets are larger than the shorter time-series."
                )

        cols = []
        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])
            cols += [model_name + n for n in model.loss.output_names]

        fcsts_df = ufp.cv_times(
            times=self.ds,
            uids=self.uids,
            indptr=self.dataset.indptr,
            h=self.h,
            test_size=test_size,
            step_size=step_size,
            id_col=id_col,
            time_col=time_col,
        )
        # the cv_times is sorted by window and then id
        fcsts_df = ufp.sort(fcsts_df, [id_col, "cutoff", time_col])

        col_idx = 0
        fcsts = np.full(
            (self.dataset.n_groups * self.h * n_windows, len(cols)),
            np.nan,
            dtype=np.float32,
        )

        for model in self.models:
            model.fit(dataset=self.dataset, val_size=val_size, test_size=test_size)
            model_fcsts = model.predict(
                self.dataset, step_size=step_size, **data_kwargs
            )

            # Append predictions in memory placeholder
            output_length = len(model.loss.output_names)
            fcsts[:, col_idx : (col_idx + output_length)] = model_fcsts
            col_idx += output_length
        # we may have allocated more space than needed
        # each serie can produce at most (serie.size - 1) // self.h CV windows
        effective_sizes = ufp.counts_by_id(fcsts_df, id_col)["counts"].to_numpy()
        needs_trim = effective_sizes.sum() != fcsts.shape[0]
        if self.scalers_ or needs_trim:
            indptr = np.arange(
                0,
                n_windows * self.h * (self.dataset.n_groups + 1),
                n_windows * self.h,
                dtype=np.int32,
            )
            if self.scalers_:
                fcsts = self._scalers_target_inverse_transform(fcsts, indptr)
            if needs_trim:
                # we keep only the effective samples of each serie from the cv results
                trimmed = np.empty_like(
                    fcsts, shape=(effective_sizes.sum(), fcsts.shape[1])
                )
                cv_indptr = np.append(0, effective_sizes).cumsum(dtype=np.int32)
                for i in range(fcsts.shape[1]):
                    ga = GroupedArray(fcsts[:, i], indptr)
                    trimmed[:, i] = ga._tails(cv_indptr)
                fcsts = trimmed

        self._fitted = True

        # Add predictions to forecasts DataFrame
        if isinstance(self.uids, pl_Series):
            fcsts = pl_DataFrame(dict(zip(cols, fcsts.T)))
        else:
            fcsts = pd.DataFrame(fcsts, columns=cols)
        fcsts_df = ufp.horizontal_concat([fcsts_df, fcsts])

        # Add original input df's y to forecasts DataFrame
        return ufp.join(
            fcsts_df,
            df[[id_col, time_col, target_col]],
            how="left",
            on=[id_col, time_col],
        )

    def cross_validation(
        self,
        df: Optional[DataFrame] = None,
        static_df: Optional[DataFrame] = None,
        n_windows: int = 1,
        step_size: int = 1,
        val_size: Optional[int] = 0,
        test_size: Optional[int] = None,
        use_init_models: bool = False,
        verbose: bool = False,
        refit: Union[bool, int] = False,
        id_col: str = "unique_id",
        time_col: str = "ds",
        target_col: str = "y",
        prediction_intervals: Optional[PredictionIntervals] = None,
        level: Optional[List[Union[int, float]]] = None,
        **data_kwargs,
    ) -> DataFrame:
        """Temporal Cross-Validation with core.NeuralForecast.

        `core.NeuralForecast`'s cross-validation efficiently fits a list of NeuralForecast
        models through multiple windows, in either chained or rolled manner.

        Parameters
        ----------
        df : pandas or polars DataFrame, optional (default=None)
            DataFrame with columns [`unique_id`, `ds`, `y`] and exogenous variables.
            If None, a previously stored dataset is required.
        static_df : pandas or polars DataFrame, optional (default=None)
            DataFrame with columns [`unique_id`] and static exogenous.
        n_windows : int (default=1)
            Number of windows used for cross validation.
        step_size : int (default=1)
            Step size between each window.
        val_size : int, optional (default=None)
            Length of validation size. If passed, set `n_windows=None`.
        test_size : int, optional (default=None)
            Length of test size. If passed, set `n_windows=None`.
        use_init_models : bool, option (default=False)
            Use initial model passed when object was instantiated.
        verbose : bool (default=False)
            Print processing steps.
        refit : bool or int (default=False)
            Retrain model for each cross validation window.
            If False, the models are trained at the beginning and then used to predict each window.
            If positive int, the models are retrained every `refit` windows.
        id_col : str (default='unique_id')
            Column that identifies each serie.
        time_col : str (default='ds')
            Column that identifies each timestep, its values can be timestamps or integers.
        target_col : str (default='y')
            Column that contains the target.
        prediction_intervals : PredictionIntervals, optional (default=None)
            Configuration to calibrate prediction intervals (Conformal Prediction).
        level : list of ints or floats, optional (default=None)
            Confidence levels between 0 and 100. Use with prediction_intervals.
        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`.
        """
        h = self.h
        if n_windows is None and test_size is None:
            raise Exception("you must define `n_windows` or `test_size`.")
        if test_size is None:
            test_size = h + step_size * (n_windows - 1)
        elif n_windows is None:
            if (test_size - h) % step_size:
                raise Exception("`test_size - h` should be module `step_size`")
            n_windows = int((test_size - h) / step_size) + 1
        else:
            raise Exception("you must define `n_windows` or `test_size` but not both")

        # Recover initial model if use_init_models.
        if use_init_models:
            self._reset_models()

        # Checks for prediction intervals
        if prediction_intervals is not None or level is not None:
            if level is None:
                warnings.warn("Level not provided, using level=[90].")
                level = [90]
            if prediction_intervals is None:
                raise Exception("You must set prediction_intervals to use level.")
            if not refit:
                raise Exception(
                    "Passing prediction_intervals and/or level is only supported with refit=True."
                )

        if not refit:

            return self._no_refit_cross_validation(
                df=df,
                static_df=static_df,
                n_windows=n_windows,
                step_size=step_size,
                val_size=val_size,
                test_size=test_size,
                verbose=verbose,
                id_col=id_col,
                time_col=time_col,
                target_col=target_col,
                **data_kwargs,
            )
        if df is None:
            raise ValueError("Must specify `df` with `refit!=False`.")
        validate_freq(df[time_col], self.freq)
        splits = ufp.backtest_splits(
            df,
            n_windows=n_windows,
            h=self.h,
            id_col=id_col,
            time_col=time_col,
            freq=self.freq,
            step_size=step_size,
            input_size=None,
        )
        results = []
        for i_window, (cutoffs, train, test) in enumerate(splits):
            should_fit = i_window == 0 or (refit > 0 and i_window % refit == 0)
            if should_fit:
                self.fit(
                    df=train,
                    static_df=static_df,
                    val_size=val_size,
                    use_init_models=False,
                    verbose=verbose,
                    id_col=id_col,
                    time_col=time_col,
                    target_col=target_col,
                    prediction_intervals=prediction_intervals,
                )
                predict_df: Optional[DataFrame] = None
            else:
                predict_df = train
            needed_futr_exog = self._get_needed_futr_exog()
            if needed_futr_exog:
                futr_df: Optional[DataFrame] = test
            else:
                futr_df = None
            preds = self.predict(
                df=predict_df,
                static_df=static_df,
                futr_df=futr_df,
                verbose=verbose,
                level=level,
                **data_kwargs,
            )
            preds = ufp.join(preds, cutoffs, on=id_col, how="left")
            fold_result = ufp.join(
                preds, test[[id_col, time_col, target_col]], on=[id_col, time_col]
            )
            results.append(fold_result)
        out = ufp.vertical_concat(results, match_categories=False)
        out = ufp.drop_index_if_pandas(out)
        # match order of cv with no refit
        first_out_cols = [id_col, time_col, "cutoff"]
        remaining_cols = [
            c for c in out.columns if c not in first_out_cols + [target_col]
        ]
        cols_order = first_out_cols + remaining_cols + [target_col]
        return ufp.sort(out[cols_order], by=[id_col, "cutoff", time_col])

    def predict_insample(self, step_size: int = 1):
        """Predict insample with core.NeuralForecast.

        `core.NeuralForecast`'s `predict_insample` uses stored fitted `models`
        to predict historic values of a time series from the stored dataframe.

        Parameters
        ----------
        step_size : int (default=1)
            Step size between each window.

        Returns
        -------
        fcsts_df : pandas.DataFrame
            DataFrame with insample predictions for all fitted `models`.
        """
        if not self._fitted:
            raise Exception(
                "The models must be fitted first with `fit` or `cross_validation`."
            )

        for model in self.models:
            if model.SAMPLING_TYPE == "recurrent":
                warnings.warn(
                    f"Predict insample might not provide accurate predictions for \
                       recurrent model {repr(model)} class yet due to scaling."
                )
                print(
                    f"WARNING: Predict insample might not provide accurate predictions for \
                      recurrent model {repr(model)} class yet due to scaling."
                )

        cols = []
        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])
            cols += [model_name + n for n in model.loss.output_names]

        # Remove test set from dataset and last dates
        test_size = self.models[0].get_test_size()

        # trim the forefront period to ensure `test_size - h` should be module `step_size
        # Note: current constraint imposes that all series lengths are equal, so we can take the first series length as sample
        series_length = self.dataset.indptr[1] - self.dataset.indptr[0]
        _, forefront_offset = np.divmod((series_length - test_size - self.h), step_size)

        if test_size > 0 or forefront_offset > 0:
            trimmed_dataset = TimeSeriesDataset.trim_dataset(
                dataset=self.dataset, right_trim=test_size, left_trim=forefront_offset
            )
            new_idxs = np.hstack(
                [
                    np.arange(
                        self.dataset.indptr[i] + forefront_offset,
                        self.dataset.indptr[i + 1] - test_size,
                    )
                    for i in range(self.dataset.n_groups)
                ]
            )
            times = self.ds[new_idxs]
        else:
            trimmed_dataset = self.dataset
            times = self.ds

        # Generate dates
        fcsts_df = _insample_times(
            times=times,
            uids=self.uids,
            indptr=trimmed_dataset.indptr,
            h=self.h,
            freq=self.freq,
            step_size=step_size,
            id_col=self.id_col,
            time_col=self.time_col,
        )

        col_idx = 0
        fcsts = np.full((len(fcsts_df), len(cols)), np.nan, dtype=np.float32)

        for model in self.models:
            # Test size is the number of periods to forecast (full size of trimmed dataset)
            model.set_test_size(test_size=trimmed_dataset.max_size)

            # Predict
            model_fcsts = model.predict(trimmed_dataset, step_size=step_size)
            # Append predictions in memory placeholder
            output_length = len(model.loss.output_names)
            fcsts[:, col_idx : (col_idx + output_length)] = model_fcsts
            col_idx += output_length
            model.set_test_size(test_size=test_size)  # Set original test_size

        # original y
        original_y = {
            self.id_col: ufp.repeat(self.uids, np.diff(self.dataset.indptr)),
            self.time_col: self.ds,
            self.target_col: self.dataset.temporal[:, 0].numpy(),
        }

        # Add predictions to forecasts DataFrame
        if isinstance(self.uids, pl_Series):
            fcsts = pl_DataFrame(dict(zip(cols, fcsts.T)))
            Y_df = pl_DataFrame(original_y)
        else:
            fcsts = pd.DataFrame(fcsts, columns=cols)
            Y_df = pd.DataFrame(original_y).reset_index(drop=True)
        fcsts_df = ufp.horizontal_concat([fcsts_df, fcsts])

        # Add original input df's y to forecasts DataFrame
        fcsts_df = ufp.join(fcsts_df, Y_df, how="left", on=[self.id_col, self.time_col])
        if self.scalers_:
            sizes = ufp.counts_by_id(fcsts_df, self.id_col)["counts"].to_numpy()
            indptr = np.append(0, sizes.cumsum())
            invert_cols = cols + [self.target_col]
            fcsts_df[invert_cols] = self._scalers_target_inverse_transform(
                fcsts_df[invert_cols].to_numpy(), indptr
            )
        return fcsts_df

    # Save list of models with pytorch lightning save_checkpoint function
    def save(
        self,
        path: str,
        model_index: Optional[List] = None,
        save_dataset: bool = True,
        overwrite: bool = False,
    ):
        """Save NeuralForecast core class.

        `core.NeuralForecast`'s method to save current status of models, dataset, and configuration.
        Note that by default the `models` are not saving training checkpoints to save disk memory,
        to get them change the individual model `**trainer_kwargs` to include `enable_checkpointing=True`.

        Parameters
        ----------
        path : str
            Directory to save current status.
        model_index : list, optional (default=None)
            List to specify which models from list of self.models to save.
        save_dataset : bool (default=True)
            Whether to save dataset or not.
        overwrite : bool (default=False)
            Whether to overwrite files or not.
        """
        # Standarize path without '/'
        if path[-1] == "/":
            path = path[:-1]

        # Model index list
        if model_index is None:
            model_index = list(range(len(self.models)))

        fs, _, _ = fsspec.get_fs_token_paths(path)
        if not fs.exists(path):
            fs.makedirs(path)
        else:
            # Check if directory is empty to protect overwriting files
            files = fs.ls(path)

            # Checking if the list is empty or not
            if files:
                if not overwrite:
                    raise Exception(
                        "Directory is not empty. Set `overwrite=True` to overwrite files."
                    )
                else:
                    fs.rm(path, recursive=True)
                    fs.mkdir(path)

        # Save models
        count_names = {"model": 0}
        alias_to_model = {}
        for i, model in enumerate(self.models):
            # Skip model if not in list
            if i not in model_index:
                continue

            model_name = repr(model)
            model_class_name = model.__class__.__name__.lower()
            alias_to_model[model_name] = model_class_name
            count_names[model_name] = count_names.get(model_name, -1) + 1
            model.save(f"{path}/{model_name}_{count_names[model_name]}.ckpt")
        with fsspec.open(f"{path}/alias_to_model.pkl", "wb") as f:
            pickle.dump(alias_to_model, f)

        # Save dataset
        if save_dataset and hasattr(self, "dataset"):
            if isinstance(self.dataset, _FilesDataset):
                raise ValueError(
                    "Cannot save distributed dataset.\n"
                    "You can set `save_dataset=False` and use the `df` argument in the predict method after loading "
                    "this model to use it for inference."
                )
            with fsspec.open(f"{path}/dataset.pkl", "wb") as f:
                pickle.dump(self.dataset, f)
        elif save_dataset:
            raise Exception(
                "You need to have a stored dataset to save it, \
                             set `save_dataset=False` to skip saving dataset."
            )

        # Save configuration and parameters
        config_dict = {
            "h": self.h,
            "freq": self.freq,
            "_fitted": self._fitted,
            "local_scaler_type": self.local_scaler_type,
            "scalers_": self.scalers_,
            "id_col": self.id_col,
            "time_col": self.time_col,
            "target_col": self.target_col,
        }
        for attr in ["prediction_intervals", "_cs_df"]:
            # conformal prediction related attributes was not available < 1.7.6
            config_dict[attr] = getattr(self, attr, None)

        if save_dataset:
            config_dict.update(
                {
                    "uids": self.uids,
                    "last_dates": self.last_dates,
                    "ds": self.ds,
                }
            )

        with fsspec.open(f"{path}/configuration.pkl", "wb") as f:
            pickle.dump(config_dict, f)

    @staticmethod
    def load(path, verbose=False, **kwargs):
        """Load NeuralForecast

        `core.NeuralForecast`'s method to load checkpoint from path.

        Parameters
        -----------
        path : str
            Directory with stored artifacts.
        kwargs
            Additional keyword arguments to be passed to the function
            `load_from_checkpoint`.

        Returns
        -------
        result : NeuralForecast
            Instantiated `NeuralForecast` class.
        """
        # Standarize path without '/'
        if path[-1] == "/":
            path = path[:-1]

        fs, _, _ = fsspec.get_fs_token_paths(path)
        files = [f.split("/")[-1] for f in fs.ls(path) if fs.isfile(f)]

        # Load models
        models_ckpt = [f for f in files if f.endswith(".ckpt")]
        if len(models_ckpt) == 0:
            raise Exception("No model found in directory.")

        if verbose:
            print(10 * "-" + " Loading models " + 10 * "-")
        models = []
        try:
            with fsspec.open(f"{path}/alias_to_model.pkl", "rb") as f:
                alias_to_model = pickle.load(f)
        except FileNotFoundError:
            alias_to_model = {}
        for model in models_ckpt:
            model_name = "_".join(model.split("_")[:-1])
            model_class_name = alias_to_model.get(model_name, model_name)
            loaded_model = MODEL_FILENAME_DICT[model_class_name].load(
                f"{path}/{model}", **kwargs
            )
            loaded_model.alias = model_name
            models.append(loaded_model)
            if verbose:
                print(f"Model {model_name} loaded.")

        if verbose:
            print(10 * "-" + " Loading dataset " + 10 * "-")
        # Load dataset
        try:
            with fsspec.open(f"{path}/dataset.pkl", "rb") as f:
                dataset = pickle.load(f)
            if verbose:
                print("Dataset loaded.")
        except FileNotFoundError:
            dataset = None
            if verbose:
                print("No dataset found in directory.")

        if verbose:
            print(10 * "-" + " Loading configuration " + 10 * "-")
        # Load configuration
        try:
            with fsspec.open(f"{path}/configuration.pkl", "rb") as f:
                config_dict = pickle.load(f)
            if verbose:
                print("Configuration loaded.")
        except FileNotFoundError:
            raise Exception("No configuration found in directory.")

        # in 1.6.4, `local_scaler_type` / `scalers_` lived on the dataset.
        # in order to preserve backwards-compatibility, we check to see if these are found on the dataset
        # in case they cannot be found in `config_dict`
        default_scalar_type = getattr(dataset, "local_scaler_type", None)
        default_scalars_ = getattr(dataset, "scalers_", None)

        # Create NeuralForecast object
        neuralforecast = NeuralForecast(
            models=models,
            freq=config_dict["freq"],
            local_scaler_type=config_dict.get("local_scaler_type", default_scalar_type),
        )

        attr_to_default = {"id_col": "unique_id", "time_col": "ds", "target_col": "y"}
        for attr, default in attr_to_default.items():
            setattr(neuralforecast, attr, config_dict.get(attr, default))
        # only restore attribute if available
        for attr in ["prediction_intervals", "_cs_df"]:
            setattr(neuralforecast, attr, config_dict.get(attr, None))

        # Dataset
        if dataset is not None:
            neuralforecast.dataset = dataset
            restore_attrs = [
                "uids",
                "last_dates",
                "ds",
            ]
            for attr in restore_attrs:
                setattr(neuralforecast, attr, config_dict[attr])

        # Fitted flag
        neuralforecast._fitted = config_dict["_fitted"]

        neuralforecast.scalers_ = config_dict.get("scalers_", default_scalars_)

        return neuralforecast

    def _conformity_scores(
        self,
        df: DataFrame,
        id_col: str,
        time_col: str,
        target_col: str,
        static_df: Optional[DataFrame],
    ) -> DataFrame:
        """Compute conformity scores.

        We need at least two cross validation errors to compute
        quantiles for prediction intervals (`n_windows=2`, specified by self.prediction_intervals).

        The exception is raised by the PredictionIntervals data class.

        df: DataFrame,
        id_col: str,
        time_col: str,
        target_col: str,
        static_df: Optional[DataFrame],
        """
        if self.prediction_intervals is None:
            raise AttributeError(
                "Please rerun the `fit` method passing a valid prediction_interval setting to compute conformity scores"
            )

        min_size = ufp.counts_by_id(df, id_col)["counts"].min()
        min_samples = self.h * self.prediction_intervals.n_windows + 1
        if min_size < min_samples:
            raise ValueError(
                "Minimum required samples in each serie for the prediction intervals "
                f"settings are: {min_samples}, shortest serie has: {min_size}. "
                "Please reduce the number of windows, horizon or remove those series."
            )

        cv_results = self.cross_validation(
            df=df,
            static_df=static_df,
            n_windows=self.prediction_intervals.n_windows,
            id_col=id_col,
            time_col=time_col,
            target_col=target_col,
        )

        kept = [time_col, id_col, "cutoff"]
        # conformity score for each model
        for model in self._get_model_names(add_level=True):
            kept.append(model)

            # compute absolute error for each model
            abs_err = abs(cv_results[model] - cv_results[target_col])
            cv_results = ufp.assign_columns(cv_results, model, abs_err)
        dropped = list(set(cv_results.columns) - set(kept))
        return ufp.drop_columns(cv_results, dropped)
