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

__all__ = ['DistributedForecast']

# Cell
from typing import Callable, Generator, List, Optional

import dask.dataframe as dd
from dask.distributed import Client, default_client

from ..core import TimeSeries, simple_predict
from ..utils import backtest_splits
from .core import DistributedTimeSeries


# Cell
class DistributedForecast:
    """Distributed pipeline encapsulation."""

    def __init__(self, model, ts: TimeSeries, client: Optional[Client] = None):
        self.model = model
        self.client = client or default_client()
        self.dts = DistributedTimeSeries(ts, self.client)
        self.model.client = self.client

    def __repr__(self) -> str:
        return f'DistributedForecast(model={self.model}, dts={self.dts})'

    def preprocess(
        self,
        data: dd.DataFrame,
        static_features: Optional[List[str]] = None,
        dropna: bool = True,
        keep_last_n: Optional[int] = None,
    ) -> dd.DataFrame:
        """Computes the transformations on each partition of `data`.

        Saves the resulting `TimeSeries` objects as well as the divisions in `data` for the forecasting step.
        Returns a dask dataframe with the computed features."""
        return self.dts.fit_transform(data, static_features, dropna, keep_last_n)

    def fit(
        self,
        data: dd.DataFrame,
        static_features: Optional[List[str]] = None,
        dropna: bool = True,
        keep_last_n: Optional[int] = None,
        **fit_kwargs,
    ) -> 'DistributedForecast':
        """Perform the preprocessing and fit the model."""
        train_ddf = self.preprocess(data, static_features, dropna, keep_last_n)
        X, y = train_ddf.drop(columns=['ds', 'y']), train_ddf.y
        self.model.fit(X, y, **fit_kwargs)
        return self

    def predict(
        self, horizon: int, predict_fn: Callable = simple_predict, **predict_fn_kwargs
    ) -> dd.DataFrame:
        """Compute the predictions for the next `horizon` steps.

        `predict_fn(model, new_x, features_order, **predict_fn_kwargs)` is called in every timestep, where:
        `model` is the trained model.
        `new_x` is a dataframe with the same format as the input plus the computed features.
        `features_order` is the list of column names that were used in the training step.
        """
        return self.dts.predict(
            self.model.model_, horizon, predict_fn, **predict_fn_kwargs
        )

    def backtest(
        self,
        data: dd.DataFrame,
        n_windows: int,
        window_size: int,
        static_features: Optional[List[str]] = None,
        dropna: bool = True,
        keep_last_n: Optional[int] = None,
        predict_fn: Callable = simple_predict,
        **predict_fn_kwargs,
    ) -> Generator[dd.DataFrame, None, None]:
        """Creates `n_windows` splits of `window_size` from `data`, trains the model
        on the training set, predicts the window and merges the actuals and the predictions
        in a dataframe.

        Returns a generator to the dataframes containing the datestamps, actual values
        and predictions."""
        for train, valid in backtest_splits(data, n_windows, window_size):
            self.fit(train, static_features, dropna, keep_last_n)
            y_pred = self.predict(window_size, predict_fn, **predict_fn_kwargs)
            y_valid = valid[['ds', 'y']]
            result = y_valid.merge(y_pred, on=['unique_id', 'ds'], how='left')
            yield result
