# %%
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")

# %%
bins = np.arange(0, 1, 0.01)
result_rdf = md.distribution.rdf(oxygen_water[-1], bins=bins)
plt.figure()
plt.plot(bins[:-1], result_rdf)
plt.xlabel(r"$r$ in nm", fontsize=16)
plt.ylabel(r"$g(r)$", fontsize=16)
plt.savefig(f"{data_dir}/rdf.png", dpi=300, bbox_inches="tight")
plt.show()
plt.close()

# %%
bins = np.arange(0, 1, 0.01)
result_rdf = md.distribution.time_average(
    partial(md.distribution.rdf, bins=bins), oxygen_water, segments=1000, skip=0.1
)
plt.figure()
plt.plot(bins[:-1], result_rdf)
plt.xlabel(r"$r$ in nm", fontsize=16)
plt.ylabel(r"$g(r)$", fontsize=16)
plt.show()
plt.close()

# %%
boron_anion = trajectory.subset(atom_name="B", residue_name="BF4")
bins = np.arange(0, 1, 0.01)
result_rdf = md.distribution.time_average(
    partial(md.distribution.rdf, bins=bins),
    oxygen_water,
    coordinates_b=boron_anion,
    segments=1000,
    skip=0.1,
)
plt.figure()
plt.plot(bins[:-1], result_rdf)
plt.xlabel(r"$r$ in nm", fontsize=16)
plt.ylabel(r"$g(r)$", fontsize=16)
plt.show()
plt.close()