import mdevaluate as md
import numpy as np

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

print(trajectory[0].atom_names)
print(trajectory[0].residue_names)

print(np.unique(trajectory[0].atom_names))
print(np.unique(trajectory[0].residue_names))

print(np.unique(trajectory[0].atom_names[trajectory[0].residue_names == "amim"]))
print(np.unique(trajectory[0].atom_names[trajectory[0].residue_names == "BF4"]))
print(np.unique(trajectory[0].atom_names[trajectory[0].residue_names == "SOL"]))

carbon_atoms = trajectory.subset(atom_name="C")
hydrogen_atoms = trajectory.subset(atom_name="H")
cation_chain_end_atoms = trajectory.subset(atom_name="CT")

cation_atoms = trajectory.subset(residue_name="amim")
anion_atoms = trajectory.subset(residue_name="BF4")
water_atoms = trajectory.subset(residue_name="SOL")

hydrogen_water = trajectory.subset(atom_name="HW", residue_name="SOL")