# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/00_LagFunctions.ipynb.

# %% auto 0
__all__ = ['lag', 'lag_uniform']

# %% ../nbs/00_LagFunctions.ipynb 2
import pandas as pd
import numpy as np
import matplotlib.cm as cm
import matplotlib.pyplot as plt

# %% ../nbs/00_LagFunctions.ipynb 6
def lag(X: pd.DataFrame, #A pandas dataframe with observations as rows and variables as columns
        lags: list, #A list with the same number of lists as the number of columns of X containing the desired lagged versions desired for the variable in that column
       )->pd.DataFrame: # A dataframe with the specified lagged variables of X

    "Return a dataframe with specific lags"
    
    res = pd.DataFrame()
    
    for x in X:
        Xlag = pd.DataFrame()
        Xlag_name = []
        Xlag_og_name = []
        for i in lags[X.columns.get_loc(x)]:
            X_shift = pd.DataFrame(X[x].shift(i))
            Xlag = pd.concat([Xlag, X_shift], axis = 1)
            
            if i == 0:
                Xlag_name.append('t')
            else:
                Xlag_name.append('t-'+f'{i}')
            Xlag_og_name.append(f'{x}')
            Xlag.columns = [Xlag_og_name,Xlag_name]
    
        res = pd.concat([res,Xlag], axis = 1)
    res = res.dropna()

    return(res)

# %% ../nbs/00_LagFunctions.ipynb 9
def lag_uniform(X: pd.DataFrame, #A pandas dataframe with observations as rows and variables as columns
        lag: int, #Integer of the maximum lag to include for all variables of X
       )->pd.DataFrame: # A dataframe with the lagged variables of X
    
    # Returns a uniformly lagged version of X

    n = np.shape(X)[0]
    p = np.shape(X)[1]
    l = lag

    res = pd.DataFrame()

    X_val = X.values

    Xlag_name = []
    Xlag_og_name = []


    for i in range(l+1):
        Xlag = pd.DataFrame(X_val[l-i:n-i])
        res = pd.concat([res,Xlag], axis = 1)
    
        if i == 0: 
            Xlag_name += ['t' for x in range(p)]
        else:
            Xlag_name += ['t-'+f'{i}' for x in range(p)]
    
        Xlag_og_name += X.columns.tolist()
    
    res.columns = [Xlag_og_name,Xlag_name]
    res.columns.names = ['Variable','Time']
    
    return res
