tests/test_md_sim_unit.py

"""
Unit tests for the MD simulation engine (no network; OpenMM only for the
thermostat test, which auto-skips when OpenMM is absent).

Covers:
- Kabsch RMSD correctness (rigid-body invariance, known displacements)
- Adaptive production length scaling (size-dependent, bounded)
- Radius of gyration (translation invariance, scaling, known geometry)
- Shrake–Ruger SASA (single atom, buried area, random-cloud bounds, rotation)
- Langevin thermostat reaches its target temperature (random data)
- Position conversion
- JSON-safe native conversion
"""

import numpy as np
import pytest

from app.tools.md_sim import (
    _kabsch_rmsd,
    _adaptive_production_steps,
    _positions_to_np,
    _radius_of_gyration,
    _sasa_shrake_ruger,
    _temperature_from_ke,
    _to_native,
    _PROBE_RADIUS_ANGSTROM,
)


def _rotz(a):
    c, s = np.cos(a), np.sin(a)
    return np.array([[c, -s, 0], [s, c, 0], [0, 0, 1]])


def _roty(a):
    c, s = np.cos(a), np.sin(a)
    return np.array([[c, 0, s], [0, 1, 0], [-s, 0, c]])


class TestKabschRMSD:
    def test_identical_coordinates(self):
        ref = np.random.RandomState(0).rand(100, 3) * 10
        assert _kabsch_rmsd(ref, ref.copy()) < 1e-9

    def test_rigid_body_invariance(self):
        """Rotation + translation must yield RMSD ~0 (Kabsch removes both)."""
        rng = np.random.RandomState(1)
        ref = rng.rand(80, 3) * 15
        R = _rotz(0.7) @ _roty(1.2) @ _rotz(0.3)
        mov = ref @ R.T + np.array([5.0, -3.0, 2.0])
        assert _kabsch_rmsd(ref, mov) < 1e-8

    def test_pure_translation_invariance(self):
        """A uniform translation alone must yield RMSD ~0."""
        rng = np.random.RandomState(7)
        ref = rng.rand(60, 3) * 10
        mov = ref + np.array([1.0, 2.0, -3.0])
        assert _kabsch_rmsd(ref, mov) < 1e-8

    def test_noise_upper_bound(self):
        """Adding noise must produce RMSD <= per-atom raw (un-aligned) RMSD, and > 0."""
        rng = np.random.RandomState(4)
        ref = rng.rand(100, 3) * 10
        noise = rng.normal(0, 0.5, ref.shape)
        mov = ref + noise
        # per-atom RMSD of un-aligned pair = sqrt(mean over ALL coords of noise^2) * sqrt(3)
        raw = float(np.sqrt((noise**2).mean())) * np.sqrt(3)
        got = _kabsch_rmsd(ref, mov)
        assert 0.0 < got <= raw * 1.001
        assert got > raw * 0.7  # optimal rotation shouldn't over-correct

    def test_single_atom(self):
        # A single atom always centers to the origin, so RMSD is 0.
        assert _kabsch_rmsd(np.array([[0.0, 0, 0]]), np.array([[1.0, 0, 0]])) < 1e-9

    def test_shape_mismatch_raises(self):
        with pytest.raises(ValueError):
            _kabsch_rmsd(np.zeros((3, 3)), np.zeros((4, 3)))

    def test_empty_input_returns_zero(self):
        assert _kabsch_rmsd(np.zeros((0, 3)), np.zeros((0, 3))) == 0.0


class TestAdaptiveProductionSteps:
    def test_small_protein_gets_target(self):
        assert _adaptive_production_steps(642) >= 100_000

    def test_larger_protein_gets_fewer_steps(self):
        steps_big = _adaptive_production_steps(30_000)
        steps_huge = _adaptive_production_steps(60_000)
        assert steps_big > steps_huge

    def test_never_exceeds_cap(self):
        assert _adaptive_production_steps(10) <= 1000 * 500  # 1 ns cap

    def test_never_below_floor(self):
        assert _adaptive_production_steps(1_000_000) >= 2 * 500

    def test_zero_atoms_returns_target(self):
        assert _adaptive_production_steps(0) == 250 * 500


class TestPositionConversion:
    def test_converts_openmm_like_positions(self):
        class P:
            def __init__(self, x, y, z):
                self.x, self.y, self.z = x, y, z

        # OpenMM positions are in nanometers; conversion must scale to Å (×10).
        positions = [P(1, 2, 3), P(4, 5, 6)]
        out = _positions_to_np(positions)
        assert out.shape == (2, 3)
        np.testing.assert_allclose(out[0], [10, 20, 30])
        np.testing.assert_allclose(out[1], [40, 50, 60])


class TestToNative:
    def test_converts_numpy_types(self):
        out = _to_native({"a": np.float32(1.5), "b": np.int64(3), "c": np.array([1.0, 2.0])})
        assert isinstance(out["a"], float)
        assert isinstance(out["b"], int)
        assert isinstance(out["c"], list)

    def test_nested_structures(self):
        out = _to_native([{"x": np.float64(1.0)}, [np.int32(2)]])
        assert isinstance(out[0]["x"], float)
        assert isinstance(out[1][0], int)


class TestRadiusOfGyration:
    def test_empty_input_zero(self):
        assert _radius_of_gyration(np.zeros((0, 3))) == 0.0

    def test_single_atom_zero(self):
        assert _radius_of_gyration(np.array([[1.0, 2.0, 3.0]])) == 0.0

    def test_known_geometry_cube_vertices(self):
        # Cube vertices at (±1, ±1, ±1): centroid at origin, every atom at distance √3.
        coords = np.array([[-1, -1, -1], [1, -1, -1], [-1, 1, -1], [-1, -1, 1],
                           [1, 1, 1], [-1, 1, 1], [1, -1, 1], [1, 1, -1]], dtype=float)
        assert abs(_radius_of_gyration(coords) - np.sqrt(3)) < 1e-9

    def test_translation_invariance(self):
        rng = np.random.default_rng(2)
        coords = rng.normal(size=(60, 3)) * 8
        shifted = coords + np.array([7.0, -4.0, 3.0])
        assert abs(_radius_of_gyration(coords) - _radius_of_gyration(shifted)) < 1e-9

    def test_scaling(self):
        rng = np.random.default_rng(5)
        coords = rng.normal(size=(40, 3)) * 5
        assert abs(_radius_of_gyration(2 * coords) - 2 * _radius_of_gyration(coords)) < 1e-9

    def test_matches_manual_formula(self):
        rng = np.random.default_rng(9)
        coords = rng.normal(size=(50, 3)) * 10
        com = coords.mean(axis=0)
        expected = float(np.sqrt(((coords - com) ** 2).sum(axis=1).mean()))
        assert abs(_radius_of_gyration(coords) - expected) < 1e-9


class TestSASA:
    def test_empty_input_zero(self):
        assert _sasa_shrake_ruger(np.zeros((0, 3)), np.zeros((0,))) == 0.0

    def test_single_atom_full_sphere(self):
        r = 1.7
        sasa = _sasa_shrake_ruger(np.array([[0.0, 0.0, 0.0]]), np.array([r]))
        expected = 4 * np.pi * (r + _PROBE_RADIUS_ANGSTROM) ** 2
        assert abs(sasa - expected) < expected * 0.02

    def test_overlap_buries_surface(self):
        # Two atoms 1 Å apart overlap heavily: surface must be between one and
        # two full spheres.
        r = 1.7
        coords = np.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]])
        radii = np.array([r, r])
        sasa = _sasa_shrake_ruger(coords, radii)
        single = 4 * np.pi * (r + _PROBE_RADIUS_ANGSTROM) ** 2
        assert single < sasa < 2 * single

    def test_separated_atoms_sum_to_twice_single(self):
        r = 1.7
        d = 2 * (r + _PROBE_RADIUS_ANGSTROM) + 1.0  # well beyond interaction range
        coords = np.array([[0.0, 0.0, 0.0], [d, 0.0, 0.0]])
        radii = np.array([r, r])
        sasa = _sasa_shrake_ruger(coords, radii)
        single = 4 * np.pi * (r + _PROBE_RADIUS_ANGSTROM) ** 2
        assert abs(sasa - 2 * single) < single * 0.05

    def test_random_cloud_bounded(self):
        rng = np.random.default_rng(11)
        coords = rng.normal(size=(300, 3)) * 12
        radii = rng.choice([1.7, 1.55, 1.52, 1.2], size=300)
        sasa = _sasa_shrake_ruger(coords, radii)
        assert sasa > 0
        sphere_sum = 4 * np.pi * ((radii + _PROBE_RADIUS_ANGSTROM) ** 2).sum()
        assert sasa < sphere_sum

    def test_rotation_invariance(self):
        rng = np.random.default_rng(13)
        coords = rng.normal(size=(200, 3)) * 10
        radii = rng.choice([1.7, 1.55, 1.52, 1.8, 1.2], size=200)
        R = _rotz(0.7) @ _roty(1.2)
        s1 = _sasa_shrake_ruger(coords, radii)
        s2 = _sasa_shrake_ruger(coords @ R.T, radii)
        assert abs(s1 - s2) / s1 < 0.02

    def test_translation_invariance(self):
        rng = np.random.default_rng(17)
        coords = rng.normal(size=(150, 3)) * 9
        radii = rng.choice([1.7, 1.55, 1.52], size=150)
        s1 = _sasa_shrake_ruger(coords, radii)
        s2 = _sasa_shrake_ruger(coords + np.array([4.0, -6.0, 2.0]), radii)
        assert abs(s1 - s2) / s1 < 0.02


class TestTemperatureFromKE:
    def test_zero_dof_returns_zero(self):
        assert _temperature_from_ke(10.0, 0) == 0.0
        assert _temperature_from_ke(10.0, -3) == 0.0

    def test_known_value(self):
        # T = 2·KE / (k_B·N_dof); solve for KE to give exactly 300 K.
        n_dof = 597
        ke = 0.5 * 0.0083144621 * 300 * n_dof
        assert abs(_temperature_from_ke(ke, n_dof) - 300.0) < 1e-9

    def test_scales_linearly_with_ke(self):
        t1 = _temperature_from_ke(10.0, 100)
        t2 = _temperature_from_ke(20.0, 100)
        assert abs(t2 - 2 * t1) < 1e-9


class TestLangevinTemperature:
    def test_random_system_reaches_thermostat_target(self):
        openmm = pytest.importorskip("openmm")
        from openmm import unit

        rng = np.random.default_rng(42)
        n_particles = 200

        # Force-free system: only the Langevin thermostat acts, so the kinetic
        # energy must converge to the Maxwell–Boltzmann average at 300 K.
        # (Adding LJ forces would make random overlapping atoms explode — a
        # physical, not numerical, effect — so we keep the test force-free.)
        system = openmm.System()
        for _ in range(n_particles):
            system.addParticle(12.0 * unit.dalton)

        integrator = openmm.LangevinMiddleIntegrator(
            300 * unit.kelvin, 1 / unit.picosecond, 2 * unit.femtoseconds)
        context = openmm.Context(system, integrator, openmm.Platform.getPlatformByName("CPU"))
        try:
            context.setPositions(rng.normal(size=(n_particles, 3)) * unit.nanometer)
            integrator.step(2000)  # warm-up from zero velocities
            temps = []
            for _ in range(50):
                integrator.step(20)
                st = context.getState(getEnergy=True)
                ke = st.getKineticEnergy().value_in_unit(unit.kilojoule_per_mole)
                temps.append(_temperature_from_ke(ke, 3 * n_particles - 3))
            mean_temp = float(np.mean(temps))
            assert 250 < mean_temp < 350, f"mean temperature {mean_temp:.1f} K far from 300 K target"
        finally:
            del context


class TestBioPythonFallback:
    """BioPython structural-analysis fallback (regression for 'str' object has
    no attribute 'name' — BioPython Atom.element is a string, not an object)."""

    _MINI_PDB = """\
ATOM      1  N   ALA A   1      11.104   6.134  -6.504  1.00 11.79           N
ATOM      2  CA  ALA A   1      11.639   6.071  -5.145  1.00 11.80           C
ATOM      3  C   ALA A   1      12.839   6.979  -4.873  1.00 11.52           C
ATOM      4  O   ALA A   1      13.419   7.636  -5.737  1.00 11.76           O
ATOM      5  CB  ALA A   1      10.514   6.523  -4.225  1.00 12.30           C
ATOM      6  N   GLY A   2      13.220   7.004  -3.610  1.00 11.67           N
ATOM      7  CA  GLY A   2      14.370   7.807  -3.226  1.00 11.95           C
ATOM      8  C   GLY A   2      14.359   8.072  -1.730  1.00 12.19           C
ATOM      9  O   GLY A   2      13.417   8.612  -1.169  1.00 12.34           O
END
"""

    def test_fallback_completes_with_sasa(self, tmp_path):
        from app.tools.md_sim import _run_biopython_analysis
        pdb_path = tmp_path / "mini.pdb"
        pdb_path.write_text(self._MINI_PDB)
        result = _run_biopython_analysis(str(pdb_path), "MINI", "minimize")
        assert result["status"] == "complete"
        assert result["engine"] == "biopython_structural"
        assert len(result["radius_of_gyration"]) >= 1
        assert result["sasa"][0]["sasa_angstrom2"] > 0
        assert result["atom_count"] > 0


class TestCTerminalOxt:
    """Regression: RCSB PDBs omit the C-terminal carboxylate oxygen (OXT).
    AMBER14's C-terminal templates require OXT while the internal template
    requires the next residue's C bond, so an unterminated C-terminus (seen on
    1TIM HIS 248) matches neither and addHydrogens() raises ValueError.
    _add_missing_terminal_oxt must add OXT so the OpenMM path succeeds instead
    of degrading to the BioPython fallback."""

    _MINI_PDB = """\
ATOM   1852  N   LYS A 247      53.278  15.624   7.791  1.00  0.00           N
ATOM   1853  CA  LYS A 247      53.240  14.342   7.088  1.00  0.00           C
ATOM   1854  C   LYS A 247      52.815  13.161   7.998  1.00  0.00           C
ATOM   1855  O   LYS A 247      52.797  13.349   9.221  1.00  0.00           O
ATOM   1856  CB  LYS A 247      54.648  14.071   6.522  1.00  0.00           C
ATOM   1857  CG  LYS A 247      55.669  13.813   7.643  1.00  0.00           C
ATOM   1858  CD  LYS A 247      57.068  13.604   7.028  1.00  0.00           C
ATOM   1859  CE  LYS A 247      58.070  13.199   8.124  1.00  0.00           C
ATOM   1860  NZ  LYS A 247      59.431  13.100   7.578  1.00  0.00           N
ATOM   1861  N   HIS A 248      52.499  12.017   7.395  1.00  0.00           N
ATOM   1862  CA  HIS A 248      52.071  10.790   8.078  1.00  0.00           C
ATOM   1863  C   HIS A 248      53.091  10.557   9.224  1.00  0.00           C
ATOM   1864  O   HIS A 248      53.300  11.600  10.100  1.00  0.00           O
ATOM   1865  CB  HIS A 248      52.029   9.501   7.220  1.00  0.00           C
ATOM   1866  CG  HIS A 248      50.801   9.422   6.366  1.00  0.00           C
ATOM   1867  ND1 HIS A 248      49.565   9.056   6.862  1.00  0.00           N
ATOM   1868  CD2 HIS A 248      50.660   9.717   5.034  1.00  0.00           C
ATOM   1869  CE1 HIS A 248      48.727   9.129   5.833  1.00  0.00           C
ATOM   1870  NE2 HIS A 248      49.338   9.522   4.722  1.00  0.00           N
END
"""

    def test_c_terminal_his_without_oxt_runs_openmm(self, tmp_path):
        pytest.importorskip("openmm")
        from app.tools.md_sim import _run_openmm
        pdb_path = tmp_path / "cterm.pdb"
        pdb_path.write_text(self._MINI_PDB)
        result = _run_openmm(str(pdb_path), "CTERM", "minimize")
        assert result["status"] == "complete"
        assert result["engine"] == "openmm"
        assert result["atom_count"] > 0
        assert result["residue_count"] == 2

    def test_add_missing_oxt_returns_zero_when_not_needed(self):
        # A structure without protein residues needs no OXT work.
        from app.tools.md_sim import _add_missing_terminal_oxt
        from openmm.app import Topology, Modeller
        from openmm import unit
        topo = Topology()
        modeller = Modeller(topo, [] * unit.nanometer)
        assert _add_missing_terminal_oxt(modeller) == 0

    def test_oxt_geometry_does_not_clash_with_sidechain(self, tmp_path):
        # Regression: the OXT position was originally computed by reflecting O
        # through C (a point reflection), which sent it straight into the
        # backbone (OXT ~1.4 A from CA/CB). That clash made the initial GBSA
        # forces enormous and tipped OpenMM into "Particle coordinate is NaN"
        # during minimizeEnergy on some platforms. OXT must be placed at the
        # ~120 deg carboxylate angle, well away from CA/CB.
        pytest.importorskip("openmm")
        from app.tools.md_sim import _add_missing_terminal_oxt
        from openmm.app import PDBFile, Modeller
        from openmm import unit
        import math
        pdb_path = tmp_path / "cterm.pdb"
        pdb_path.write_text(self._MINI_PDB)
        pdb = PDBFile(str(pdb_path))
        modeller = Modeller(pdb.topology, pdb.positions)
        assert _add_missing_terminal_oxt(modeller) == 1

        pos = {}
        for atom in modeller.topology.atoms():
            pos[atom.name] = modeller.positions[atom.index].value_in_unit(unit.nanometer)
        def dist(a, b):
            return math.dist(pos[a], pos[b])
        # Bond to C preserved, and no clash with backbone/sidechain atoms.
        assert 0.10 < dist("C", "OXT") < 0.16
        assert dist("CA", "OXT") > 0.20
        assert dist("CB", "OXT") > 0.20
        assert dist("O", "OXT") > 0.20
        # Near-planar carboxylate, ~120 deg O-C-OXT angle (not 180 deg).
        c, o, oxt = pos["C"], pos["O"], pos["OXT"]
        v1 = (o[0]-c[0], o[1]-c[1], o[2]-c[2])
        v2 = (oxt[0]-c[0], oxt[1]-c[1], oxt[2]-c[2])
        n1 = math.dist(o, c); n2 = math.dist(oxt, c)
        cosang = (v1[0]*v2[0] + v1[1]*v2[1] + v1[2]*v2[2]) / (n1*n2)
        angle = math.degrees(math.acos(max(-1.0, min(1.0, cosang))))
        assert 60.0 < angle < 180.0