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

__all__ = ['__single', '__double', 'holt_winters', '__smoothing', 'SINGLE', 'DOUBLE']

# Cell
import numpy as np
from numba import jit
import pandas as pd

from time_series_model_basics import simulate_data

# Cell
def __single(
    ts,
    alpha=0.1,
):

    if alpha > 1 or alpha < 0:
        raise Exception("Alpha should be on the interval (0,1)")

    n = len(ts)
    f = np.full((n,), np.nan)

    f[1] = ts[0]

    for i in range(2, n):
        f[i] = alpha * ts[i - 1] + (1 - alpha) * f[i - 1]

    return f

# Cell
@jit(nopython=True)
def __double(
    ts,
    alpha=0.1,
    beta=0.1,
    s_0=None,
    b_0=None,
    m=1,
):

    if alpha > 1 or alpha < 0:
        raise Exception("alpha should be on the interval (0,1)")
    if beta > 1 or beta < 0:
        raise Exception("beta should be on the interval (0,1)")

    n = len(ts) + m - 1
    s = np.full((n,), np.nan)
    b = np.full((n,), np.nan)
    f = np.full((n,), np.nan)
    tss = np.append(ts, np.full(m, np.nan))

    s[0] = tss[0] if s_0 is None else s_0
    b[0] = tss[1] - tss[0] if b_0 is None else b_0

    for i in range(1, n):

        s[i] = alpha * tss[i] + (1 - alpha) * (s[i - 1] + b[i - 1])
        b[i] = beta * (s[i] - s[i - 1]) + (1 - beta) * b[i - 1]

        f[i - 1 + m] = s[i - 1] + m * b[i - 1]

    return f

# Cell
@jit(nopython=True)
def holt_winters(
    ts,
    l=1,
    alpha=0.1,
    beta=0.1,
    gamma=0.1,
    s_0=None,
    b_0=None,
    m=1,
):

    if alpha > 1 or alpha < 0:
        raise Exception("alpha should be on the interval (0,1)")
    if beta > 1 or beta < 0:
        raise Exception("beta should be on the interval (0,1)")
    if gamma > 1 or gamma < 0:
        raise Exception("gamma should be on the interval (0,1)")

    if (l <= 0) and (l - int(l) != 0):
        raise Exception("l should be a Natural Number")

    n = len(ts) + m
    #complete cycles
    N = (n - m) // l

    if N == 0:
        raise Exception("it should be at least one complete cycle")

    s = np.full((n,), np.nan)
    b = np.full((n,), np.nan)
    f = np.full((n,), np.nan)
    c = np.full((n,), np.nan)
    A = np.full((N,), np.nan)

    tss = np.append(ts, np.full(m, np.nan))

    s[0] = tss[0] if s_0 is None else s_0

    b[0] = (1 / l**2) * np.sum(tss[l:2 * l] - tss[:l])

    for j in range(1, N + 1):
        A[j - 1] = (1 / l) * np.sum(tss[l * (j - 1):l * (j - 1) + l])

    for i in range(l):
        c[i] = (1 / N) * np.sum(tss[i:l * (N - 1) + i + 1:l] / A)

    for i in range(1, n):
        s[i] = alpha * (tss[i] - c[i - l]) + (1 - alpha) * (s[i - 1] + b[i - 1])
        b[i] = beta * (s[i] - s[i - 1]) + (1 - beta) * b[i - 1]
        c[i] = gamma * (tss[i] - s[i - 1] - b[i - 1]) + (1 - gamma) * c[i - l]

        f[i - 1 + m] = s[i - i] + m * b[i] + c[i - l + 1 + ((m - 1) % l)]

    return f

# Cell
def __smoothing(
    *args,
    kind=None,
    df=None,
    ts_col=None,
    **kwargs,
):
    """
    Forescasts  with Exponential Smoothing

    -----
    Parameters
    -----

    kind : str , default 'simple'. It can be:

         'simple' : Single Smoothing.
         'double' : Double Smoothing.
    df : DataFrame, default None. If df is None a dataframe
         with simulated data is generated.

    ts_col : str, default None. The column name with the
         time series. Set to 'time_series' when df is None.
    *args : int.  Each value represent a Moving
           Average Forecast and its corresponding
           window size
    **kwargs : passed to __pandas_time_series

    -------
    Returns
    -------

    tuple: First element is the Pandas Dataframe and the second
           is ploty's figure object
    """

    if len(args) == 0:
        raise Exception("list in *args should include at least one value")

    if kind not in ('single', 'double'):
        raise Exception("Only supported methods are: 'simple', 'weighted'")

    if df is None:
        df = simulate_data.pandas_time_series(**kwargs)
        ts_col = 'time_series'

    if ts_col is None:
        raise Exception('When df is not None, ts_col should be specified')

    ts = df[ts_col].to_numpy()
    cols = [ts_col]
    for a in args:
        col = f'{kind}_{a}'
        cols.append(col)
        if kind == 'single':
            df[col] = __single(
                ts=ts,
                alpha=a,
            )
        elif kind == 'double':
            df[col] = __double(
                ts=ts,
                alpha=a[0],
                beta=a[1],
            )

    fig = df[cols].plot(
        backend='plotly',
        title=f'Smoothing {kind}',
    )

    fig.update_layout(template="plotly_dark",)

    return df, fig

# Cell
def SINGLE(
    *args,
    df=None,
    ts_col=None,
    **kwargs,
):
    """
    Forescasts  with Single Smoothing

    -----
    Parameters
    -----
    df : DataFrame, default None. If df is None a dataframe
         with simulated data is generated.
    ts_col : str, default 'time_series'. The column name with the
         time series. Ignored if df is None.
    *args : int.  Each value represent a Single Smoothing
         parameter
    **kwargs : passed to __pandas_time_series
    -------
    Returns
    -------
    tuple: First element is the Pandas Dataframe and the second
           is ploty's fig object
    """

    return __smoothing(
        *args,
        kind='single',
        df=df,
        ts_col=ts_col,
        **kwargs,
    )

# Cell
def DOUBLE(*args, df=None, ts_col=None, **kwargs):
    """
    Forescasts  with Double Smoothing

    -----
    Parameters
    -----
    df : DataFrame, default None. If df is None a dataframe
         with simulated data is generated.
    ts_col : str, default 'time_series'. The column name with the
         time series. Ignored if df is None.
    *args : Each pair reprents the two parameters
         of the double smoothing
    **kwargs : passed to __pandas_time_series
    -------
    Returns
    -------
    tuple: First element is the Pandas Dataframe and the second
           is ploty's fig object
    """

    return __smoothing(
        *args,
        kind='double',
        df=df,
        ts_col=ts_col,
        **kwargs,
    )