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

atom_indices_a = trajectory.subset(
    atom_name="HW.", residue_name="SOL"
).atom_subset.indices
atom_indices_b = trajectory.subset(
    atom_name="OW", residue_name="SOL"
).atom_subset.indices
atom_indices_b = np.vstack([atom_indices_b] * 2).T.reshape((-1,))
OH_vectors = md.coordinates.vectors(
    trajectory,
    atom_indices_a=atom_indices_a,
    atom_indices_b=atom_indices_b,
    normed=True,
)
print(OH_vectors[-1])

atom_indices_a = trajectory.subset(
    atom_name="B", residue_name="BF4"
).atom_subset.indices
atom_indices_b = trajectory.subset(
    atom_name="F.", residue_name="BF4"
).atom_subset.indices
atom_indices_a = np.vstack([atom_indices_a] * 4).T.reshape((-1,))
BF_vectors = md.coordinates.vectors(
    trajectory,
    atom_indices_a=atom_indices_b,
    atom_indices_b=atom_indices_a,
    normed=True,
)
print(BF_vectors[-1])

atom_indices = trajectory.subset(residue_name="SOL").atom_subset.indices
water_dipole_vectors = md.coordinates.dipole_vector(
    trajectory, atom_indices=atom_indices, normed=True
)
print(water_dipole_vectors[-1])

atom_indices_a = trajectory.subset(
    atom_name="NA1", residue_name="amim"
).atom_subset.indices
atom_indices_b = trajectory.subset(
    atom_name="CR", residue_name="amim"
).atom_subset.indices
atom_indices_c = trajectory.subset(
    atom_name="CW1", residue_name="amim"
).atom_subset.indices
cation_ring_normal = md.coordinates.normal_vectors(
    trajectory,
    atom_indices_a=atom_indices_a,
    atom_indices_b=atom_indices_b,
    atom_indices_c=atom_indices_c,
    normed=True,
)
print(cation_ring_normal[-1])