from numpy import *
import numpy as np
import math
import os
import matplotlib.pyplot as plt
from io import StringIO

def file_to_array_csv(filename):
    dir = os.path.dirname(__file__)
    with open(os.path.join(dir, filename), 'r') as file:
        return np.array([line.strip().split(',') for line in file]).astype(np.float)

def file_to_array_txt(filename):
    dir = os.path.dirname(__file__)
    data = np.loadtxt(os.path.join(dir, filename))
    return data

def corr(sinal_A,sinal_B,janela):
    # sinalA e sinalB tem que ter o mesmo temanho
    # e mesma frequência de aquisição
    # Janela: se vc quer -100 < n < 100, janela = 100, janela em PONTOS
    N = len(sinal_A)
    mean_A = sinal_A.sum(axis=0)/N
    mean_B = sinal_B.sum(axis=0)/N
    flut_A = sinal_A - mean_A * np.ones((1, N))
    flut_B = sinal_B - mean_B * np.ones((1, N))
    var_A = np.dot(flut_A,flut_A.T)/N
    var_B = np.dot(flut_B, flut_B.T) / N
    sigma_A = math.sqrt(var_A)
    sigma_B = math.sqrt(var_B)
    c = np.zeros((1, 2*janela+1))
    n_array = np.array([range(-janela, janela + 1)])
    for n in range(-janela, janela + 1):
        M = N - abs(n)
        term = 0
        for m in range(M):
            if n < 0:
                term = term + flut_A[0,m]*flut_B[0,m+abs(n)]/(M*sigma_A*sigma_B)
            if n >= 0:
                term = term + flut_A[0,m+abs(n)]*flut_B[0,m]/(M*sigma_A*sigma_B)
            # print(term)
        c[0,janela+n] = term
        # print(c)
    return c,n_array

if(__name__ == '__main__'):

    filename = 'TCABR.txt'
    A_orig = file_to_array_txt(filename)
    sinal_t = A_orig[:,0]
    sinal_A = A_orig[:,1]
    sinal_B = A_orig[:,2]
    janela = 25
    c, n = corr(sinal_A,sinal_B,janela)
    plt.plot(n[0,:],c[0,:])
    plt.title('Correlação cruzada TCABR_A e TCABR_B')
    plt.xlabel('Defasagem (n)')
    plt.ylabel('Correlação')
    plt.grid(True)
    plt.show()