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

__all__ = ['random_time_series', '__simple', '__double', 'holt_winters', 'SIMPLE', 'DOUBLE']

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

# Cell
def random_time_series(
    num_points=30,
    level=3,
    trend_amplitude=1 / 3,
    sinus_amplitude=1,
    noise_amplitude=2 / 1,
    start=0,
    finish=4 * np.pi,
):

    interval = np.linspace(
        start,
        finish,
        num=num_points + 1,
    )

    season = np.sin(interval[:-1])
    trend = np.vectorize(lambda x: trend_amplitude * x)(interval[:-1])
    noise = noise_amplitude * np.random.random((num_points,))
    signal = level + season + trend + noise
    return interval, signal

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

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

    n = len(ts) + 1
    f = np.zeros((n,))

    f[0] = ts[0]

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

    f[n-1] = f[n - 2]
    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
    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)]
        print(i-1)

    print(l, n, m)
    return f

# Cell
def SIMPLE(*args):
    interval, d = random_time_series()

    time_series = {'time_series': np.append(d, [np.nan])}

    preds = {f'alpha_{alpha}': __simple(
        ts=d,
        alpha=alpha,
    ) for alpha in args}

    time_series.update(preds)

    df = pd.DataFrame(
        time_series,
        index=interval,
    )
    df.index = interval

    fig = df.plot(
        backend='plotly',
        title=f'Simple Smoothing Average',
    )

    fig.update_layout(template="plotly_dark",)

    return df, fig

# Cell
def DOUBLE(*args):

    interval, d = random_time_series(num_points=35)

    time_series = {'time_series': np.append(d, [np.nan])}

    preds = {
        f'{ab}': __double(
            ts=d,
            alpha=ab[0],
            beta=ab[1],
            m=1,
        ) for ab in args
    }

    time_series.update(preds)

    df = pd.DataFrame(
        time_series,
        index=interval,
    )
    df.index = interval

    fig = df.plot(
        backend='plotly',
        title=f'Double Smooth. Forecast one time-steps ahead',
    )

    fig.update_layout(template="plotly_dark",)

    return df, fig