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

__all__ = ['AutoARIMAProphet']

# Cell
import sys
from copy import deepcopy

import numpy as np
import statsmodels.api as sm
import pandas as pd
from ..arima import AutoARIMA

if sys.version_info.minor != 6 or (sys.platform not in ['win32', 'cygwin']):
    try:
        from prophet import Prophet
    except ModuleNotFoundError as e:
        msg = (
            '{e}. To use prophet adapters you have to install '
            'prophet. Please run `pip install prophet`. '
            'Note that it is recommended to install prophet '
            'using conda environments due to dependencies.'
        )
        raise ModuleNotFoundError(msg) from e
elif sys.version_info.minor == 6 and (sys.platform in ['win32', 'cygwin']):
    try:
        from fbprophet import Prophet
    except ModuleNotFoundError as e:
        msg = (
            '{e}. To use prophet adapters you have to install '
            'fbprophet. Please run `pip install fbprophet`. '
            'Note that it is recommended to install prophet '
            'using conda environments due to dependencies.'
        )
        raise ModuleNotFoundError(msg) from e


# Cell
class AutoARIMAProphet(Prophet):
    """Returns best ARIMA model using external variables created
    by the Prophet interface.

    This class receives as parameters the same as prophet.Prophet
    and statsforecast.arima.AutoARIMA.

    If your pipeline uses Prophet you can simply replace Prophet
    with AutoARIMAProphet and you'll be using AutoARIMA instead
    of Prophet.
    """

    def __init__(self,
                 growth='linear',
                 changepoints=None,
                 n_changepoints=25,
                 changepoint_range=0.8,
                 yearly_seasonality='auto',
                 weekly_seasonality='auto',
                 daily_seasonality='auto',
                 holidays=None,
                 seasonality_mode='additive',
                 seasonality_prior_scale=10.0,
                 holidays_prior_scale=10.0,
                 changepoint_prior_scale=0.05,
                 mcmc_samples=0,
                 interval_width=0.80,
                 uncertainty_samples=1000,
                 stan_backend=None,
                 d=None,
                 D=None,
                 max_p=5,
                 max_q=5,
                 max_P= 2,
                 max_Q= 2,
                 max_order= 5,
                 max_d=2,
                 max_D=1,
                 start_p=2,
                 start_q=2,
                 start_P=1,
                 start_Q=1,
                 stationary=False,
                 seasonal=True,
                 ic='aicc',
                 stepwise=True,
                 nmodels=94,
                 trace=False,
                 approximation=False,
                 method=None,
                 truncate=None,
                 test='kpss',
                 test_kwargs=None,
                 seasonal_test='seas',
                 seasonal_test_kwargs=None,
                 allowdrift=False,
                 allowmean=False,
                 blambda=None,
                 biasadj=False,
                 parallel=False,
                 num_cores=2,
                 period=1):
        Prophet.__init__(self,
                         growth,
                         changepoints,
                         n_changepoints,
                         changepoint_range,
                         yearly_seasonality,
                         weekly_seasonality,
                         daily_seasonality,
                         holidays,
                         seasonality_mode,
                         seasonality_prior_scale,
                         holidays_prior_scale,
                         changepoint_prior_scale,
                         mcmc_samples,
                         interval_width,
                         uncertainty_samples,
                         stan_backend)
        self.arima = AutoARIMA(d=d,
                               D=D,
                               max_p=max_p,
                               max_q=max_q,
                               max_P=max_P,
                               max_Q=max_Q,
                               max_order=max_order,
                               max_d=max_d,
                               max_D=max_D,
                               start_p=start_p,
                               start_q=start_q,
                               start_P=start_P,
                               start_Q=start_Q,
                               stationary=stationary,
                               seasonal=seasonal,
                               ic=ic,
                               stepwise=stepwise,
                               nmodels=nmodels,
                               trace=trace,
                               approximation=approximation,
                               method=method,
                               truncate=truncate,
                               test=test,
                               test_kwargs=test_kwargs,
                               seasonal_test=seasonal_test,
                               seasonal_test_kwargs=seasonal_test_kwargs,
                               allowdrift=allowdrift,
                               allowmean=allowmean,
                               blambda=blambda,
                               biasadj=biasadj,
                               parallel=parallel,
                               num_cores=num_cores,
                               period=period)

    def fit(self, df, disable_seasonal_features=True, **kwargs):
        """Fit the AutoARIMAProphet model.

        Parameters
        ----------
        df: pd.DataFrame containing the history. Must have columns ds (date
            type) and y, the time series.
        disable_seasonal_features: bool
            Wheter disable seasonal features generated by Prophet.
        kwargs: Additional arguments.

        Returns
        -------
        The AutoARIMAProphet object.
        """
        if self.history is not None:
            raise Exception('Prophet object can only be fit once. '
                            'Instantiate a new object.')
        if ('ds' not in df) or ('y' not in df):
            raise ValueError(
                'Dataframe must have columns "ds" and "y" with the dates and '
                'values respectively.'
            )
        history = df[df['y'].notnull()].copy()
        if history.shape[0] < 2:
            raise ValueError('Dataframe has less than 2 non-NaN rows.')
        self.history_dates = pd.to_datetime(pd.Series(df['ds'].unique(), name='ds')).sort_values()

        history = self.setup_dataframe(history, initialize_scales=True)
        self.history = history
        self.set_auto_seasonalities()
        seasonal_features, prior_scales, component_cols, modes = (
            self.make_all_seasonality_features(history))
        self.train_component_cols = component_cols
        self.component_modes = modes
        self.fit_kwargs = deepcopy(kwargs)
        if disable_seasonal_features:
            seas = tuple(self.seasonalities.keys())
            seasonal_features = seasonal_features.loc[:,~seasonal_features.columns.str.startswith(seas)]
        self.xreg_cols = seasonal_features.columns

        y = history['y'].values
        X = seasonal_features.values if not seasonal_features.empty else None
        self.arima = self.arima.fit(y=y, X=X)

        return self

    def predict(self, df=None):
        """Predict using the AutoARIMAProphet model.

        Parameters
        ----------
        df: pd.DataFrame with dates for predictions (column ds.
            If not provided, predictions are
            made on the history.

        Returns
        -------
        A pd.DataFrame with the forecast components.
        """
        if self.history is None:
            raise Exception('Model has not been fit.')

        if df is None:
            df = self.history.copy()
        else:
            if df.shape[0] == 0:
                raise ValueError('Dataframe has no rows.')
            df = self.setup_dataframe(df.copy())

        seasonal_features = self.make_all_seasonality_features(df)[0].loc[:, self.xreg_cols]

        ds_forecast = set(df['ds'])
        h = len(ds_forecast - set(self.history['ds']))
        if h > 0:
            X = seasonal_features.values[-h:] if not seasonal_features.empty else None
            df2 = self.arima.predict(h=h,
                                     X=X,
                                     level=int(100*self.interval_width))
        else:
            df2 = pd.DataFrame()
        if len(ds_forecast) > h:
            in_sample = self.arima.predict_in_sample(level=int(100*self.interval_width))
            df2 = pd.concat([in_sample, df2]).reset_index(drop=True)

        yhat = df2.pop('mean')
        df2.columns = ['yhat_lower', 'yhat_upper']
        df2.insert(0, 'yhat', yhat)
        df2.insert(0, 'ds', df['ds'])

        return df2