# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/01_Correlation_Matrices.ipynb.

# %% auto 0
__all__ = ['imposter_matrix', 'correlation_matrix', 'ACM', 'partial_correlation_matrix', 'PACM']

# %% ../nbs/01_Correlation_Matrices.ipynb 2
import pandas as pd
import numpy as np
import matplotlib.cm as cm
import matplotlib.pyplot as plt
from lag import *

# %% ../nbs/01_Correlation_Matrices.ipynb 6
def imposter_matrix(X: pd.DataFrame, # Matrix to be randomized
                   random_state = 42, # Random state to be used
                   )-> pd.DataFrame: # The input matrix with shuffled values for each column
    d = pd.DataFrame(0, index=np.arange(len(X)), columns=X.columns)
    
    for column in d:
        d[column] = X[column].sample(frac = 1, random_state=random_state).reset_index(drop=True)
        
    return d
    

# %% ../nbs/01_Correlation_Matrices.ipynb 8
def correlation_matrix(X: pd.DataFrame, # Matrix 1 for calculating the correlation matrix
                       y: pd.DataFrame, # Matrix 2 for calculating the correlation matrix
                      )-> np.ndarray: # Correlation matrix as a Numpy array
    n = np.shape(X)[0]
    
    X_scaled = (X-np.mean(X,axis = 0))/np.std(X, axis = 0, ddof = 1)
    y_scaled = (y-np.mean(y,axis = 0))/np.std(y, axis = 0, ddof = 1)
    
    correlation_matrix = np.dot(y_scaled.T,X_scaled) / (n-1)
    
    return correlation_matrix
    

# %% ../nbs/01_Correlation_Matrices.ipynb 10
def ACM(X: pd.DataFrame, # Dataframe of raw data for which to calculate the ACM
        lag: int, # Lag to calculate correlation against
        )-> pd.DataFrame: # Autocorrelation Matrix (ACM)
    
    n = np.shape(X)[0] # Number of observations
    p = np.shape(X)[1] # Number of variables
    
    X_val = X.values
    l = lag
    
    
    Xtrim = X_val[:n-l] # Observations from start to n-lag
    Xlag =  X_val[l:]   # Observations from lag to end
    
    # Calculate correlation matrix between Xtrim and Xlag
    ACM = correlation_matrix(Xtrim, Xlag)
    
    # New column names
    t0 = ['t' for i in range(p)]
    columns = [X.columns.to_list(),t0]
    
    # New index names
    t = ['t-'+f'{lag}' for i in range(p)]
    index = [X.columns.to_list(),t]
    
    ACM = pd.DataFrame(ACM)
    ACM.index = index
    ACM.columns = columns
    
    return ACM

# %% ../nbs/01_Correlation_Matrices.ipynb 12
def partial_correlation_matrix(X: pd.DataFrame, # Matrix for calculating partial correlation
                       )-> np.ndarray: # Partial correlation matrix as a Numpy array 
    V = X.cov()
    Vi = np.linalg.pinv(V)
    D = np.diag(np.sqrt(1/np.diag(Vi)))
    pcor = -1 * (D @ Vi @ D)
    pcor[np.diag_indices_from(pcor)] = 1
    
    return pcor

# %% ../nbs/01_Correlation_Matrices.ipynb 14
def PACM(X: pd.DataFrame, # Matrix for calculating partial autocorrelation
         lag: int, # Lag to calculate partial correlation against
        ):
    
    p = np.shape(X)[1]
    
    X_lag = lag_uniform(X, lag)
    
    pcor_mat = partial_correlation_matrix(X_lag)
    
    pcor_df = pd.DataFrame(pcor_mat)
    pcor_df.columns = X_lag.columns
    pcor_df.index = X_lag.columns
    
    
    PACM = pcor_df.iloc[lag*p:(lag*p+p),0:p]
    
    return PACM
