# %%
from functools import partial

import numpy as np
import matplotlib.pyplot as plt

import mdevaluate as md

data_dir = "/data/skloth/python_packages/mdevaluate_examples/plots"
path_to_sim = "/data/skloth/sim/IL_water/C4MIM_BF4_x50/T300_nvt"
trajectory = md.open(path_to_sim, topology="run.tpr", trajectory="out/traj_full.xtc")

oxygen_water = trajectory.subset(atom_name="OW", residue_name="SOL")

# %%
print(md.correlation.msd(trajectory[0], trajectory[-1]))

result_msd_pbc = md.correlation.shifted_correlation(
    md.correlation.msd, oxygen_water, segments=100, skip=0.1, average=True
)
result_msd_nojump = md.correlation.shifted_correlation(
    md.correlation.msd, oxygen_water.nojump, segments=100, skip=0.1, average=True
)

plt.figure()
plt.plot(result_msd_pbc[0], result_msd_pbc[1], ".", label="pbc")
plt.plot(result_msd_nojump[0], result_msd_nojump[1], ".", label="nojump")
plt.legend()
plt.xscale("log")
plt.yscale("log")
plt.xlabel(r"$t$ in ps", fontsize=16)
plt.ylabel(r"$\langle r^2\rangle(t)$ in nm$^2$", fontsize=16)
plt.show()
plt.close()

# %%
cation = trajectory.subset(residue_name="amim")
cation_indices = cation.atom_subset.indices
cation_com = md.coordinates.center_of_masses(trajectory, atom_indices=cation_indices)

result_msd_cation = md.correlation.shifted_correlation(
    md.correlation.msd, cation.nojump, segments=100, skip=0.1, average=True
)
result_msd_cation_com = md.correlation.shifted_correlation(
    md.correlation.msd, cation_com.nojump, segments=100, skip=0.1, average=True
)

plt.figure()
plt.plot(result_msd_cation[0], result_msd_cation[1], ".", label="cation all atoms")
plt.plot(result_msd_cation_com[0], result_msd_cation_com[1], ".", label="cation com")
plt.legend()
plt.xscale("log")
plt.yscale("log")
plt.xlabel(r"$t$ in ps", fontsize=16)
plt.ylabel(r"$\langle r^2\rangle(t)$ in nm$^2$", fontsize=16)
plt.show()
plt.close()