import pandas as pd
import numpy as np
from scipy.stats import pearsonr, spearmanr
from scipy.signal import correlate
import numpy as np
import statsmodels as st
from st.tsa.stattools import grangercausalitytests
from sklearn.metrics import mutual_info_score


def calculate_correlation(data, method='pearson'):
    if method == 'pearson':
        corr_matrix = pearsonr(data)
    elif method == 'spearman':
        corr_matrix = spearmanr(data)
    else:
        raise ValueError("Invalid correlation method. Must be 'pearson' or 'spearman'.")
    correlations = pd.DataFrame(corr_matrix[0], columns=data.columns[1:], index=data.columns[1:]) 
    return correlations


def calculate_cross_correlation(data1, data2):
    xcorr = correlate(data1, data2)
    lags = np.arange(-len(data1) + 1, len(data1))
    return xcorr, lags

def calculate_phase_sync(data1, data2):
    phase_diff = np.angle(np.exp(1j * (np.angle(np.fft.fft(data1)) - np.angle(np.fft.fft(data2)))))
    return np.abs(np.mean(np.exp(1j * phase_diff)))



def calculate_granger_causality(x, y, maxlag=10):
    data = np.column_stack((x, y))
    results = grangercausalitytests(data, maxlag=maxlag, verbose=False)
    p_value = results[maxlag][0]['ssr_ftest'][1]
    return p_value

def calculate_mutual_information(x, y):

    mi = mutual_info_score(x, y)
    return mi
