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

__all__ = ['ses', 'ses_optimized', 'adida', 'historic_average', 'croston_classic', 'croston_sba', 'croston_optimized',
           'seasonal_window_average', 'imapa', 'naive', 'seasonal_naive', 'random_walk_with_drift', 'window_average',
           'seasonal_exponential_smoothing', 'seasonal_ses_optimized', 'tsb', 'auto_arima', 'ets']

# Cell
from itertools import count
from numbers import Number
from typing import Collection, List, Optional, Sequence, Tuple

import numpy as np
import pandas as pd
from numba import njit
from scipy.optimize import minimize

from .arima import auto_arima_f, forecast_arima, fitted_arima
from .ets import ets_f, forecast_ets

# Internal Cell
@njit
def _ses_fcst_mse(x: np.ndarray, alpha: float) -> Tuple[float, float]:
    """Perform simple exponential smoothing on a series.

    This function returns the one step ahead prediction
    as well as the mean squared error of the fit.
    """
    smoothed = x[0]
    n = x.size
    mse = 0.
    fitted = np.full(n, np.nan, np.float32)

    for i in range(1, n):
        smoothed = (alpha * x[i - 1] + (1 - alpha) * smoothed).item()
        error = x[i] - smoothed
        mse += error * error
        fitted[i] = smoothed

    mse /= n
    forecast = alpha * x[-1] + (1 - alpha) * smoothed
    return forecast, mse, fitted


def _ses_mse(alpha: float, x: np.ndarray) -> float:
    """Compute the mean squared error of a simple exponential smoothing fit."""
    _, mse, _ = _ses_fcst_mse(x, alpha)
    return mse


@njit
def _ses_forecast(x: np.ndarray, alpha: float) -> float:
    """One step ahead forecast with simple exponential smoothing."""
    forecast, _, fitted = _ses_fcst_mse(x, alpha)
    return forecast, fitted


@njit
def _demand(x: np.ndarray) -> np.ndarray:
    """Extract the positive elements of a vector."""
    return x[x > 0]


@njit
def _intervals(x: np.ndarray) -> np.ndarray:
    """Compute the intervals between non zero elements of a vector."""
    y = []

    ctr = 1
    for val in x:
        if val == 0:
            ctr += 1
        else:
            y.append(ctr)
            ctr = 1

    y = np.array(y)
    return y


@njit
def _probability(x: np.ndarray) -> np.ndarray:
    """Compute the element probabilities of being non zero."""
    return (x != 0).astype(np.int32)


def _optimized_ses_forecast(
        x: np.ndarray,
        bounds: Sequence[Tuple[float, float]] = [(0.1, 0.3)]
    ) -> float:
    """Searches for the optimal alpha and computes SES one step forecast."""
    alpha = minimize(
        fun=_ses_mse,
        x0=(0,),
        args=(x,),
        bounds=bounds,
        method='L-BFGS-B'
    ).x[0]
    forecast, fitted = _ses_forecast(x, alpha)
    return forecast, fitted


@njit
def _chunk_sums(array: np.ndarray, chunk_size: int) -> np.ndarray:
    """Splits an array into chunks and returns the sum of each chunk."""
    n = array.size
    n_chunks = n // chunk_size
    sums = np.empty(n_chunks)
    for i, start in enumerate(range(0, n, chunk_size)):
        sums[i] = array[start : start + chunk_size].sum()
    return sums

# Cell
@njit
def ses(X, h, future_xreg, fitted, alpha):
    y = X[:, 0] if X.ndim == 2 else X
    fcst, _, fitted_vals = _ses_fcst_mse(y, alpha)
    mean = np.full(h, fcst, np.float32)
    fcst = {'mean': mean}
    if fitted:
        fcst['fitted'] = fitted_vals
    return fcst

# Cell
def ses_optimized(X, h, future_xreg, fitted):
    y = X[:, 0] if X.ndim == 2 else X
    fcst, res = _optimized_ses_forecast(y, [(0.01, 0.99)])
    mean = np.full(h, fcst, np.float32)
    fcst = {'mean': mean}
    if fitted:
        fcst['fitted'] = fitted
    return fcst

# Cell
def adida(X, h, future_xreg, fitted):
    if fitted:
        raise NotImplementedError('return fitted')
    y = X[:, 0] if X.ndim == 2 else X
    if (y == 0).all():
        return {'mean': np.repeat(np.float32(0), h)}
    y_intervals = _intervals(y)
    mean_interval = y_intervals.mean()
    aggregation_level = round(mean_interval)
    lost_remainder_data = len(y) % aggregation_level
    y_cut = y[lost_remainder_data:]
    aggregation_sums = _chunk_sums(y_cut, aggregation_level)
    sums_forecast, _ = _optimized_ses_forecast(aggregation_sums)
    forecast = sums_forecast / aggregation_level
    mean = np.repeat(forecast, h)
    return {'mean': mean}

# Cell
@njit
def historic_average(X, h, future_xreg, fitted):
    y = X[:, 0] if X.ndim == 2 else X
    mean = np.repeat(y.mean(), h)
    fcst = {'mean': mean}
    if fitted:
        fitted_vals = np.full(y.size, np.nan, y.dtype)
        fitted_vals[1:] = y.cumsum()[:-1] / np.arange(1, y.size)
        fcst['fitted'] = fitted_vals
    return fcst

# Cell
@njit
def croston_classic(X, h, future_xreg, fitted):
    if fitted:
        raise NotImplementedError('return fitted')
    y = X[:, 0] if X.ndim == 2 else X
    yd = _demand(y)
    yi = _intervals(y)
    ydp, _ = _ses_forecast(yd, 0.1)
    yip, _ = _ses_forecast(yi, 0.1)
    mean = ydp / yip
    return {'mean': mean}


@njit
def croston_sba(X, h, future_xreg, fitted):
    if fitted:
        raise NotImplementedError('return fitted')
    y = X[:, 0] if X.ndim == 2 else X
    mean = croston_classic(y, h, future_xreg, fitted)
    mean['mean'] *= 0.95
    return mean


def croston_optimized(X, h, future_xreg, fitted):
    if fitted:
        raise NotImplementedError('return fitted')
    y = X[:, 0] if X.ndim == 2 else X
    yd = _demand(y)
    yi = _intervals(y)
    ydp, _ = _optimized_ses_forecast(yd)
    yip, _ = _optimized_ses_forecast(yi)
    mean = ydp / yip
    return {'mean': mean}


@njit
def seasonal_window_average(
    X: np.ndarray,
    h: int,
    future_xreg,
    fitted,
    season_length: int,
    window_size: int,
) -> np.ndarray:
    if fitted:
        raise NotImplementedError('return fitted')
    y = X[:, 0] if X.ndim == 2 else X
    min_samples = season_length * window_size
    if y.size < min_samples:
        return {'mean': np.full(h, np.nan, np.float32)}
    season_avgs = np.zeros(season_length, np.float32)
    for i, value in enumerate(y[-min_samples:]):
        season = i % season_length
        season_avgs[season] += value / window_size
    out = np.empty(h, np.float32)
    for i in range(h):
        out[i] = season_avgs[i % season_length]
    return {'mean': out}

def imapa(X, h, future_xreg, fitted):
    if fitted:
        raise NotImplementedError('return fitted')
    y = X[:, 0] if X.ndim == 2 else X
    if (y == 0).all():
        return {'mean': np.repeat(np.float32(0), h)}
    y_intervals = _intervals(y)
    mean_interval = y_intervals.mean().item()
    max_aggregation_level = round(mean_interval)
    forecasts = np.empty(max_aggregation_level, np.float32)
    for aggregation_level in range(1, max_aggregation_level + 1):
        lost_remainder_data = len(y) % aggregation_level
        y_cut = y[lost_remainder_data:]
        aggregation_sums = _chunk_sums(y_cut, aggregation_level)
        forecast, _ = _optimized_ses_forecast(aggregation_sums)
        forecasts[aggregation_level - 1] = (forecast / aggregation_level)
    forecast = forecasts.mean()
    mean = np.repeat(forecast, h)
    return {'mean': mean}


@njit
def naive(X, h, future_xreg, fitted):
    y = X[:, 0] if X.ndim == 2 else X
    mean = np.repeat(y[-1], h).astype(np.float32)
    if fitted:
        fitted_vals = np.full(y.size, np.nan, np.float32)
        fitted_vals[1:] = np.roll(y, 1)[1:]
        return {'mean': mean, 'fitted': fitted_vals}
    return {'mean': mean}

# Cell
@njit
def seasonal_naive(X, h, future_xreg, fitted, season_length):
    y = X[:, 0] if X.ndim == 2 else X
    if y.size < season_length:
        return {'mean': np.full(h, np.nan, np.float32)}
    season_vals = np.empty(season_length, np.float32)
    fitted_vals = np.full(y.size, np.nan, np.float32)
    for i in range(season_length):
        s_naive = naive(y[i::season_length], 1, None, fitted)
        season_vals[i] = s_naive['mean'].item()
        if fitted:
            fitted_vals[i::season_length] = s_naive['fitted']
    out = np.empty(h, np.float32)
    for i in range(h):
        out[i] = season_vals[i % season_length]
    fcst = {'mean': out}
    if fitted:
        fcst['fitted'] = fitted_vals
    return fcst

# Cell
@njit
def random_walk_with_drift(X, h, future_xreg, fitted):
    y = X[:, 0] if X.ndim == 2 else X
    slope = (y[-1] - y[0]) / (y.size - 1)
    mean = slope * (1 + np.arange(h)) + y[-1]
    fcst = {'mean': mean.astype(np.float32)}
    if fitted:
        fitted_vals = np.full(y.size, np.nan, dtype=np.float32)
        fitted_vals[1:] = (slope + y[:-1]).astype(np.float32)
        fcst['fitted'] = fitted_vals
    return fcst

# Cell
@njit
def window_average(X, h, future_xreg, fitted, window_size):
    if fitted:
        raise NotImplementedError('return fitted')
    y = X[:, 0] if X.ndim == 2 else X
    if y.size < window_size:
        return {'mean': np.full(h, np.nan, np.float32)}
    wavg = y[-window_size:].mean()
    mean = np.repeat(wavg, h)
    return {'mean': mean}

# Cell
@njit
def seasonal_exponential_smoothing(X, h, future_xreg, fitted, season_length, alpha):
    y = X[:, 0] if X.ndim == 2 else X
    if y.size < season_length:
        return {'mean': np.full(h, np.nan, np.float32)}
    season_vals = np.empty(season_length, np.float32)
    fitted_vals = np.full(y.size, np.nan, np.float32)
    for i in range(season_length):
        season_vals[i], fitted_vals[i::season_length] = _ses_forecast(y[i::season_length], alpha)
    out = np.empty(h, np.float32)
    for i in range(h):
        out[i] = season_vals[i % season_length]
    fcst = {'mean': out}
    if fitted:
        fcst['fitted'] = fitted_vals
    return fcst

# Cell
def seasonal_ses_optimized(X, h, future_xreg, fitted, season_length):
    y = X[:, 0] if X.ndim == 2 else X
    if y.size < season_length:
        return {'mean': np.full(h, np.nan, np.float32)}
    season_vals = np.empty(season_length, np.float32)
    fitted_vals = np.full(y.size, np.nan, np.float32)
    for i in range(season_length):
        season_vals[i], fitted_vals[i::season_length] = _optimized_ses_forecast(y[i::season_length], [(0.01, 0.99)])
    out = np.empty(h, np.float32)
    for i in range(h):
        out[i] = season_vals[i % season_length]
    fcst = {'mean': out}
    if fitted:
        fcst['fitted'] = fitted_vals
    return fcst

# Cell
@njit
def tsb(X, h, future_xreg, fitted, alpha_d, alpha_p):
    if fitted:
        raise NotImplementedError('return fitted')
    y = X[:, 0] if X.ndim == 2 else X
    if (y == 0).all():
        return {'mean': np.repeat(np.float32(0), h)}
    yd = _demand(y)
    yp = _probability(y)
    ypf, _ = _ses_forecast(yp, alpha_p)
    ydf, _ = _ses_forecast(yd, alpha_d)
    forecast = np.float32(ypf * ydf)
    mean = np.repeat(forecast, h)
    return {'mean': mean}

# Cell
def auto_arima(X: np.ndarray, h: int, future_xreg=None, fitted: bool = False, season_length: int = 1,
               approximation: bool = False, level: Optional[Tuple[int]] = None) -> np.ndarray:
    y = X[:, 0] if X.ndim == 2 else X
    xreg = X[:, 1:] if (X.ndim == 2 and X.shape[1] > 1) else None
    with np.errstate(invalid='ignore'):
        mod = auto_arima_f(
            y,
            xreg=xreg,
            period=season_length,
            approximation=approximation,
            allowmean=False, allowdrift=False #not implemented yet
        )
    fcst = forecast_arima(mod, h, xreg=future_xreg, level=level)
    mean = fcst['mean']
    if fitted:
        return {'mean': mean, 'fitted': fitted_arima(mod)}
    if level is None:
        return {'mean': mean}
    return {
        'mean': mean,
        **{f'lo-{l}': fcst['lower'][f'{l}%'] for l in reversed(level)},
        **{f'hi-{l}': fcst['upper'][f'{l}%'] for l in level},
    }

# Cell
def ets(X: np.ndarray, h: int, future_xreg=None, fitted: bool = False,
        season_length: int = 1,
        model: str = 'ZZZ') -> np.ndarray:
    y = X[:, 0] if X.ndim == 2 else X
    xreg = X[:, 1:] if (X.ndim == 2 and X.shape[1] > 1) else None
    mod = ets_f(y, m=season_length, model=model)
    fcst = forecast_ets(mod, h)
    keys = ['mean']
    if fitted:
        keys.append('fitted')
    return {key: fcst[key] for key in keys}