# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/modules/05_transformations.ipynb.

# %% ../nbs/modules/05_transformations.ipynb #ae4862ca
from __future__ import annotations
from typing import List, Dict, Optional, Callable, Tuple, Any, Union

from pyexpat import model
import pandas as pd
import numpy as np
from numba import jit
from scipy.stats import boxcox
from scipy.special import inv_boxcox
from numba import jit
import matplotlib.pyplot as plt
import warnings
warnings.filterwarnings("ignore")

#------------------------------------------------------------------------------
# Box-Cox Transformation Utility Functions
#------------------------------------------------------------------------------

def box_cox_transform(x,
                      shift: bool = False,
                      box_cox_lmda: float = None
                      ):
    """
    Applies a Box-Cox transformation to a series x.

    Parameters
    ----------
    x :array-like:
        The input data to be transformed.
    shift :bool
        Whether to shift the data by 1 before transformation to handle zeros.
    box_cox_lmda :float or None:
        The lambda parameter for the Box-Cox transformation. If None, it will be estimated from the data.

    Returns
    -------
    transformed_data and lmbda:
        The Box-Cox transformed data and the lambda used for transformation.
    """
    if (box_cox_lmda == None):
        if shift ==True:
            transformed_data, lmbda = boxcox((np.array(x)+1))
        else:
            transformed_data, lmbda = boxcox(np.array(x))
            
    if (box_cox_lmda != None):
        if shift ==True:
            lmbda = box_cox_lmda
            transformed_data = boxcox((np.array(x)+1), lmbda)
        else:
            lmbda = box_cox_lmda
            transformed_data = boxcox(np.array(x), lmbda)
    return transformed_data, lmbda

# %% auto #0
__all__ = ['box_cox_transform', 'back_box_cox_transform', 'fourier_terms', 'rolling_mean', 'rolling_quantile', 'rolling_std',
           'expanding_mean', 'expanding_std', 'expanding_quantile']

# %% ../nbs/modules/05_transformations.ipynb #18d295f7
def back_box_cox_transform(y_pred: np.ndarray,
                           lmda: float,
                           shift: bool = False,
                           box_cox_biasadj: bool = False
                           ) -> np.ndarray:
    """
    Inverse Box-Cox transform.

    Parameters
    ----------
    y_pred : array-like
        The Box-Cox transformed forecast to be back-transformed.
    lmda : float
        The lambda parameter used in the Box-Cox transformation.
    shift : bool
        Whether the original data was shifted by 1 before transformation.
    box_cox_biasadj : bool
        Whether to apply bias adjustment to the back-transformed forecast.
    
    Returns
    -------
    forecast : array-like
        The back-transformed forecast.
    """
    if (box_cox_biasadj==False):
        if shift == True:
            forecast = inv_boxcox(y_pred, lmda)-1
        else:
            forecast = inv_boxcox(y_pred, lmda)
            
    if (box_cox_biasadj== True):
        pred_var = np.var(y_pred)
        if shift == True:
            if lmda ==0:
                forecast = np.exp(y_pred)*(1+pred_var/2)-1
            else:
                forecast = ((lmda*y_pred+1)**(1/lmda))*(1+((1-lmda)*pred_var)/(2*((lmda*y_pred+1)**2)))-1
        else:
            if lmda ==0:
                forecast = np.exp(y_pred)*(1+pred_var/2)
            else:
                forecast = ((lmda*y_pred+1)**(1/lmda))*(1+((1-lmda)*pred_var)/(2*((lmda*y_pred+1)**2)))
    return forecast

# %% ../nbs/modules/05_transformations.ipynb #f0b192bc
#------------------------------------------------------------------------------
# Transformation Utility Functions
#------------------------------------------------------------------------------

def fourier_terms(
    start_end_index: tuple,
    period: int,
    num_terms: int,
    frequency: str = "D"
) -> pd.DataFrame:
    """
    Generate Fourier terms for a given time index or range of indices.

    Parameters
    ----------
    start_end_index : tuple
        A tuple containing the start and end indices (can be integers or datetime strings).
    period : int
        The period of the seasonality (e.g., 7 for weekly, 12 for monthly).
    num_terms : int
        The number of Fourier terms to generate.
    frequency : str, optional
        The frequency of the time index (default is "D" for daily). For example, "D" for daily, "M" for monthly, "H" for hourly, etc. It is used when the start and end indices are datetime strings to create the appropriate date range.
    
    Returns
    -------
    pd.DataFrame
        A DataFrame containing the Fourier terms with the appropriate index.
    """
    start, end = start_end_index

    # Decide index type
    if isinstance(start, (int, np.integer)) and isinstance(end, (int, np.integer)):
        index = pd.RangeIndex(start, end+1)  # exclusive of end
        t = np.arange(len(index))
    else:
        # Convert to datetime if string given
        start, end = pd.to_datetime(start), pd.to_datetime(end)
        index = pd.date_range(start, end, freq=frequency)
        t = np.arange(len(index))

    # Build Fourier terms
    terms = {
        f'sin_{k}_{period}': np.sin(2 * np.pi * k * t / period)
        for k in range(1, num_terms + 1)
    }
    terms.update({
        f'cos_{k}_{period}': np.cos(2 * np.pi * k * t / period)
        for k in range(1, num_terms + 1)
    })

    return pd.DataFrame(terms, index=index)


# %% ../nbs/modules/05_transformations.ipynb #aaca54e0
class rolling_mean:
    def __init__(self,
                 window_size: int,
                 shift: int = 1,
                 min_samples: int = 1
                 ):

        """
        
        A class to compute the rolling mean of a time series with specified window size and shift.

        Parameters
        ----------
        window_size : int
            The size of the rolling window.
        shift : int, optional
            The number of periods to shift the data before applying the rolling mean (default is 1).
        min_samples : int, optional
            The minimum number of observations in the window required to have a value (default is 1).

        Returns
        -------
        None

        """
        self.shift = shift
        self.window_size = window_size
        self.min_samples = min_samples

    def __call__(self,
                 data: Union[pd.Series, np.ndarray],
                 is_forecast: bool = False
                 ) -> pd.Series:
        """
        Compute the rolling mean of the input data.
        
        Parameters
        ----------
        data : array-like
            The input time series data for which to compute the rolling mean.
        is_forecast : bool, optional
            Whether the data is a forecast (default is False). If True, the shift will be adjusted to align with the forecasted period. For example, if it's a forecast, for the forecasting next value, we might want to shift by one less than usual to align with the forecasted period

        Returns
        -------
        pd.Series
            A Series containing the rolling mean of the input data, shifted and computed according to the specified parameters.
        """

        if is_forecast:
            # For example, if it's a forecast, for the forecasting next value, we might want to shift by one less than usual to align with the forecasted period
            return pd.Series(data).shift(self.shift-1).rolling(self.window_size, min_periods=self.min_samples).mean()
        # If not a forecast, apply the usual shift
        else:
            # If not a forecast, apply the usual shift
            return pd.Series(data).shift(self.shift).rolling(self.window_size, min_periods=self.min_samples).mean()
        
    def get_name(self):
        return f"rolling_mean_{self.window_size}_{self.shift}"

# %% ../nbs/modules/05_transformations.ipynb #206604d3
class rolling_quantile:
    def __init__(self,
                 window_size: int,
                 quantile: float,
                 shift: int = 1,
                 min_samples: int = 1
                 ):

        """
        A class to compute the rolling quantile of a time series with specified window size, quantile, and shift.

        Parameters
        ----------
        window_size : int
            The size of the rolling window.
        quantile : float
            The quantile to compute (between 0 and 1).
        shift : int, optional
            The number of periods to shift the data before applying the rolling quantile (default is 1).
        min_samples : int, optional
            The minimum number of observations in the window required to have a value (default is 1).
        
        Returns
        -------
        None
        """
        self.shift = shift
        self.window_size = window_size
        self.quantile = quantile
        self.min_samples = min_samples

    def __call__(self,
                 data: Union[pd.Series, np.ndarray],
                 is_forecast: bool = False
                 ) -> pd.Series:
        """
        
        Compute the rolling quantile of the input data.
        
        Parameters:
        ----------
        data : array-like
            The input time series data for which to compute the rolling quantile.
        is_forecast : bool, optional
            Whether the data is a forecast (default is False). If True, the shift will be adjusted to align with the forecasted period. For example, if it's a forecast, for the forecasting next value, we might want to shift by one less than usual to align with the forecasted period
        
        Returns
        -------
        pd.Series
            A Series containing the rolling quantile of the input data, shifted and computed according to the specified parameters.
        """
        if is_forecast:
            return pd.Series(data).shift(self.shift-1).rolling(self.window_size, min_periods=self.min_samples).quantile(self.quantile)
        # If not a forecast, apply the usual shift
        else:
            # If not a forecast, apply the usual shift
            return pd.Series(data).shift(self.shift).rolling(self.window_size, min_periods=self.min_samples).quantile(self.quantile)
    def get_name(self):
        return f"rolling_quantile_{self.window_size}_{self.quantile}_{self.shift}"

# %% ../nbs/modules/05_transformations.ipynb #92824459
class rolling_std:

    def __init__(self,
                 window_size: int,
                 shift: int = 1,
                 min_samples: int = 1):

        """
        A class to compute the rolling standard deviation of a time series with specified window size and shift.

        Parameters
        ----------
        window_size : int
            The size of the rolling window.
        shift : int, optional
            The number of periods to shift the data before applying the rolling standard deviation (default is 1).
        min_samples : int, optional
            The minimum number of observations in the window required to have a value (default is 1).
        
        Returns
        -------
        None

        """

        self.shift = shift
        self.window_size = window_size
        self.min_samples = min_samples

    def __call__(self,
                 data: Union[pd.Series, np.ndarray],
                 is_forecast: bool = False
                 ) -> pd.Series:
        """
        Compute the rolling standard deviation of the input data.
        
        Parameters
        ----------
        data : array-like
            The input time series data for which to compute the rolling standard deviation.
        is_forecast : bool, optional
            Whether the data is a forecast (default is False). If True, the shift will be adjusted to align with the forecasted period. For example, if it's a forecast, for the forecasting next value, we might want to shift by one less than usual to align with the forecasted period
        
        Returns
        -------
        pd.Series
            A Series containing the rolling standard deviation of the input data, shifted and computed according to the specified parameters.
        """
        # Return the rolling std with the specified window size and minimum samples
        if is_forecast:
            return pd.Series(data).shift(self.shift-1).rolling(self.window_size, min_periods=self.min_samples).std()
        else:
            return pd.Series(data).shift(self.shift).rolling(self.window_size, min_periods=self.min_samples).std()
    def get_name(self):
        return f"rolling_std_{self.window_size}_{self.shift}"   

# %% ../nbs/modules/05_transformations.ipynb #ced41853
class expanding_mean:
    def __init__(self,
                 shift: int = 1
                 ):
        """
        A class to compute the expanding mean of a time series with specified shift.

        Parameters
        ----------
        shift: int
            The number of periods to shift time series data.
        
        Returns
        -------
        None
        """
        
        self.shift = shift
    def __call__(self, data, is_forecast=False):
        if is_forecast:
            return pd.Series(data).shift(self.shift-1).expanding().mean()
        else:
            return pd.Series(data).shift(self.shift).expanding().mean()
    def get_name(self):
        return f"expanding_mean_{self.shift}"

# %% ../nbs/modules/05_transformations.ipynb #ec1137bb
class expanding_std:
    def __init__(self,
                 shift: int = 1):
        """
        A class to compute the expanding standard deviation of a time series with specified shift.
        
        Parameters
        ----------
        shift : int, optional
            The number of periods to shift the data before applying the expanding standard deviation (default is 1).
        
        Returns
        -------
        None
        """
        self.shift = shift
    def __call__(self, data, is_forecast=False):
        if is_forecast:
            return pd.Series(data).shift(self.shift-1).expanding().std()
        else:
            return pd.Series(data).shift(self.shift).expanding().std()
    def get_name(self):
        return f"expanding_std_{self.shift}"

# %% ../nbs/modules/05_transformations.ipynb #8a1068f3
class expanding_quantile:
    def __init__(self,
                 shift: int = 1,
                 quantile: float = 0.5):
        
        """
        A class to compute the expanding quantile of a time series with specified shift and quantile.
        
        Parameters
        ----------
        shift : int, optional
            The number of periods to shift the data before applying the expanding quantile (default is 1).
        quantile : float, optional
            The quantile to compute (between 0 and 1) (default is 0.5).
        
        Returns
        -------
        None
        """
        self.shift = shift
        self.quantile = quantile

    def __call__(self, data,
                 is_forecast=False):
        """
        Compute the expanding quantile.
        """
        if is_forecast:
            return pd.Series(data).shift(self.shift-1).expanding().quantile(self.quantile)
        else:
            return pd.Series(data).shift(self.shift).expanding().quantile(self.quantile)
    def get_name(self):
        return f"expanding_quantile_{self.quantile}_{self.shift}"   
        
