r"""
A bent triatomic molecule: coupled bond-stretch and angle-bend vibrations
=============================================================================

:class:`~chemistrykit.md.systems.pair_potentials.HarmonicMolecule` builds a
small molecular-mechanics "force field" out of nothing but harmonic bond
(:class:`~chemistrykit.md.systems.pair_potentials.HarmonicBond`) and
harmonic angle-bend (:class:`~chemistrykit.md.systems.pair_potentials.HarmonicAngle`)
terms -- the classical, spring-and-hinge picture of molecular structure
that founded molecular mechanics. A bent, water-like triatomic (one
heavy vertex atom, two light outer atoms, equilibrium angle
104.5 degrees) is displaced away from its equilibrium geometry and left
to vibrate under velocity-Verlet integration alone (no thermostat): both
bond lengths and the bond angle oscillate indefinitely about their
equilibrium values, and the total energy is conserved to within
integration error.
"""

# %%
import matplotlib.pyplot as plt
import numpy as np

from chemistrykit.md.systems.pair_potentials import HarmonicMolecule

theta0 = np.deg2rad(104.5)  # a water-like equilibrium bond angle
r0 = 1.0
m_heavy, m_light = 16.0, 1.0

# Equilibrium geometry: the vertex (heavy) atom at the origin, the two
# light atoms placed symmetrically about it.
positions = np.array(
    [
        [r0 * np.sin(theta0 / 2), r0 * np.cos(theta0 / 2), 0.0],  # light atom 0
        [0.0, 0.0, 0.0],  # heavy vertex atom 1
        [-r0 * np.sin(theta0 / 2), r0 * np.cos(theta0 / 2), 0.0],  # light atom 2
    ]
)
velocities = np.zeros_like(positions)
positions[0] += np.array([0.15, 0.05, 0.0])  # displace: stretch + open the angle

masses = [m_light, m_heavy, m_light]
bonds = [(0, 1, 300.0, r0), (1, 2, 300.0, r0)]
angles = [(0, 1, 2, 50.0, theta0)]
molecule = HarmonicMolecule(positions, velocities, masses, bonds=bonds, angles=angles)

dt, n_steps = 0.001, 4000
t = np.arange(n_steps + 1) * dt
bond01, bond12, angle_deg, total_energy = [], [], [], []


def _record():
    p = molecule.positions
    r01 = float(np.linalg.norm(p[0] - p[1]))
    r12 = float(np.linalg.norm(p[2] - p[1]))
    v1, v2 = p[0] - p[1], p[2] - p[1]
    theta = float(np.arccos(np.dot(v1, v2) / (np.linalg.norm(v1) * np.linalg.norm(v2))))
    bond01.append(r01)
    bond12.append(r12)
    angle_deg.append(np.degrees(theta))
    total_energy.append(molecule.kinetic_energy() + molecule.potential_energy())


_record()
for _ in range(n_steps):
    molecule.step(dt)
    _record()

fig, axes = plt.subplots(1, 2, figsize=(11, 4.5))
axes[0].plot(t, bond01, label="bond 0-1 (stretched)")
axes[0].plot(t, bond12, label="bond 1-2")
axes[0].axhline(r0, color="gray", linestyle=":", linewidth=0.8, label="equilibrium r0")
axes[0].set_xlabel("t")
axes[0].set_ylabel("bond length")
axes[0].set_title("Bond-stretch vibration")
axes[0].legend(fontsize=8)

axes[1].plot(t, angle_deg, color="darkorange")
axes[1].axhline(np.degrees(theta0), color="gray", linestyle=":", linewidth=0.8, label="equilibrium theta0")
axes[1].set_xlabel("t")
axes[1].set_ylabel("bond angle (degrees)")
axes[1].set_title("Angle-bend vibration")
axes[1].legend(fontsize=8)
fig.tight_layout()

# %%
# Both bond lengths and the bond angle oscillate indefinitely about their
# equilibrium values -- coupled through the shared vertex atom -- and the
# total energy stays constant to within the integrator's usual small,
# bounded error, exactly the symplectic-integration behavior expected of
# velocity-Verlet applied to a conservative force field:

total_energy = np.asarray(total_energy)
drift = abs(total_energy.max() - total_energy.min()) / abs(total_energy[0])
print(f"Relative energy drift over the run: {drift:.2e}")

plt.show()