from functools import partial

import matplotlib.pyplot as plt
from matplotlib import cm

import mdevaluate as md

data_dir = "/data/skloth/python_packages/mdevaluate_examples/plots"
path_to_sim = "/data/skloth/sim/silica_pore/tip4p2005/D3_L6_S4.9_R0/T300_isochor/"
trajectory = md.open(path_to_sim, topology="run.tpr", trajectory="out/traj_full.xtc")

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

time, result_all = md.correlation.shifted_correlation(
    partial(md.correlation.isf, q=22.7),
    oxygen_water,
    segments=10,
    skip=0.1
    )

plt.figure()
plt.plot(time, result_all)
plt.xscale("log")
plt.savefig(f"{data_dir}/selector.png", dpi=300, bbox_inches="tight")