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

__all__ = ['random_time_series', 'simple', 'weighted', 'SMA', 'WMA']

# 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
@jit(nopython=True)
def simple(
    ts,
    n=1,
):

    l = len(ts)
    if n > l:
        raise Exception('n should be less than the length of the time-series')

    f = np.zeros((l + 1,))
    f[:n] = np.full(n, np.nan)

    for i in range(l - n + 1):
        f[n + i] = np.mean(ts[i:n + i])
    return f

# Cell
@jit(nopython=True)
def weighted(
    ts,
    ws,
):

    l = len(ts)
    n = len(ws)
    w = np.sum(ws)

    if n > l:
        raise Exception('The length of ws should be less than the length of ts')

    f = np.zeros((l + 1,))
    f[:n] = np.full(n, np.nan)

    for i in range(l - n + 1):
        f[n + i] = (1 / w) * np.dot(ws, ts[i:n + i])
    return f

# Cell
def SMA(*args):
    """
    Creates a Dataframe with simulated data and forescasts
    with a set of Simple Moving Averages

    -----
    Parameters
    -----

    args: list of integers. Each value represent a Moving
          Average Forecast and its corresponding
          window size

    -------
    Returns
    -------
    tuple: First element is the Pandas Dataframe and the second
           is ploty's fig object
    """

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



    interval, d = random_time_series()

    preds = {f'ma_{n}': simple(ts=d, n=n) for n in args}
    ts = {'time_series': np.append(d, np.nan)}
    ts.update(preds)
    df = pd.DataFrame(
       ts,
        index=interval,
    )
    df.index = interval
    df.index.name = 'Time Series Units'

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

    fig.update_layout(template="plotly_dark",)

    return df, fig

# Cell
def WMA(*args):
    """
    Creates a Dataframe with simulated data and forescasts
    with a set of Simple Moving Averages

    -----
    Parameters
    -----

    args: list of lists. Each list represents a set
          of Weights for the WMA

    -------
    Returns
    -------
    tuple: First element is the Pandas Dataframe and the second
           is ploty's fig object
    """

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

    interval, d = random_time_series()

    preds = {
        f'wma_{ws}': weighted(
            ts=d,
            ws=np.array(ws, dtype=float),
        ) for ws in args
    }
    ts = {'time_series': np.append(d, np.nan)}
    ts.update(preds)
    df = pd.DataFrame(
        ts,
        index=interval,
    )
    df.index = interval
    df.index.name = 'Time Series Units'

    fig = df.plot(
        backend='plotly',
        title=f'Weighted Moving Average',
    )

    fig.update_layout(template="plotly_dark",)

    return df, fig