import numpy as np
import matplotlib.pyplot as plt
from tqdm import tqdm

def normalize_cm(matriz):
    '''
    Traslada el sistema al centro de masas en cada tiempo.
    '''
    tiempo = matriz.shape[0]
    matriz2 = np.zeros_like(matriz)
    for i in range(tiempo):
        for j in range(3):
            matriz2[i,:,j] = matriz[i,:,j] - np.mean(matriz[i,:,j])
    return matriz2

def calcular_msd(datanorm):

    dt_max = datanorm.shape[0] // 2
    msd = np.zeros(dt_max - 1)
    tiempo = np.arange(1, dt_max) * 100 # 100 es el paso de tiempo
    
    '''
    Nota procedimiento:
        - Calculamos tf - ti (siempre es la resta entre mismas particulas porque no se toca las otras dim)
        Ex. dt = 1, entonces datanorm[dt:] va desde 1 hasta final, mientras que
            datanorm[:-dt] va desde el comienzo hasta el penultimo (resta segundo con primero, y asi hasta final menos penultimo)
            Las dimensiones tienen que cambiar porque si dt = 1, tenemos 1599 en vez de 1600 y el CRITERIO es hasta la mitad
        - Calculamos la norma al cuadrado
        --> Tenemos x valores de deltat para cada particula, hacemos el promedio que sera el promedio del sistema para ese dt
    '''
    
    for dt in tqdm(range(1, dt_max), desc="Calculando MSD"):
        displac = datanorm[dt:] - datanorm[:-dt] 
        dist2 = np.sum(displac ** 2, axis=2)
        msd[dt - 1] = np.mean(dist2)
    
    return msd, tiempo

# %% CARGA DE DATOS / NORMALIZACION AL CM

data1 = np.load("matrizcart_xi2_t06.npy")  
data2 = np.load("matrizcartb_xi2_t06.npy")
data3 = np.load("matrizcartc_xi2_t06.npy")

data = np.concatenate((data1, data2, data3))
datanorm = normalize_cm(data)

# %% CÁLCULO DEL MSD 

msd, tiempo = calcular_msd(datanorm)
np.save('msddef_xi2_t06', msd) # guardar

# %% GRAFICA

xp1 = np.log(tiempo)[0] # ajuste desde el inicio
yp1 = np.log(msd)[0]

x1 = np.linspace(xp1, xp1 + 6, 100)
y1= yp1 + 1 * (x1 - xp1)

plt.figure()
plt.plot(np.log(tiempo), np.log(msd), label="MSD")
plt.plot(x1, y1, "--", label="Pendiente = 1", color="orange")
plt.xlabel("log(Δt)")
plt.ylabel("log(MSD)")
plt.title("MSD del sistema (Xi = 2, kbT = 0.6)", fontsize=14, fontweight="bold")
plt.grid(True)
plt.legend()
plt.tight_layout()

plt.savefig("difusiondef_xi2_t06.png", dpi=300, bbox_inches='tight')  # guardar

plt.show()