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

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

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

from time_series_model_basics import simulate_data

# 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.full((l,), 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,
):
    if type(ws) is list:
        ws = np.array(ws, dtype=float)

    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.full((l,), np.nan)
    #f[:n] = np.full(n, np.nan)

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

# Cell
def __moving_average(
    *args,
    kind='simple',
    df=None,
    ts_col=None,
    **kwargs,
):
    """
    Forecasts  with a Moving Average

    -----
    Parameters
    -----

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

         'simple' : Simple Moving Average.
         'weighted' : Weighted Moving Average.
    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 ('simple', 'weighted'):
        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 == 'simple':
            df[col] = simple(ts=ts, n=a)
        elif kind == 'weighted':
            df[col] = weighted(ts=ts, ws=a)

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

    fig.update_layout(template="plotly_dark",)

    return df, fig

# Cell
def SMA(*args, df=None, ts_col=None, **kwargs):
    """
    Forecasts  with Simple Moving Average

    -----
    Parameters
    -----

    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 fig object
    """

    return __moving_average(*args,
                            kind='simple',
                            df=df,
                            ts_col=ts_col,
                            **kwargs)

# Cell
def WMA(*args, df=None, ts_col='time_series', **kwargs):
    """
    Forecasts with Weighted Moving Average

    -----
    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 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 fig object
    """

    return __moving_average(*args,
                            kind='weighted',
                            df=df,
                            ts_col=ts_col,
                            **kwargs)