bio_files 0.5.1

//! [Docs](https://manual.gromacs.org/current/onlinehelp/gmx-solvate.html)
//! [More docs](https://manual.gromacs.org/current/how-to/topology.html#water-solvation)

use std::fmt::Write as _;

use lin_alg::f32::Vec3;

use crate::gromacs::SOLVENT_GRO_NAME;

/// C+P from `dynamics`.
#[derive(Clone, Debug, PartialEq)]
pub struct WaterInitTemplate {
    // velocity is o velocity, instead of 3 separate velocities
    pub o_posits: Vec<Vec3>,
    pub h0_posits: Vec<Vec3>,
    pub h1_posits: Vec<Vec3>,
    pub o_velocities: Vec<Vec3>,
    pub h0_velocities: Vec<Vec3>,
    pub h1_velocities: Vec<Vec3>,
    // todo: Cache these, or infer?
    /// This must correspond to the positions. Cached.
    pub bounds: (Vec3, Vec3),
}

impl WaterInitTemplate {
    /// Create `.gro` file text for use as the `gmx solvate -cs` template.
    ///
    /// Writes 4 atoms per molecule (OW, HW1, HW2, MW) in GROMACS `.gro` format.
    /// The MW virtual-site position is computed from the OW/HW1/HW2 positions
    /// using the OPC `[virtual_sites3]` coefficients (a = b = 0.147803).
    ///
    /// Positions are converted Å → nm. Velocities (Å/ps → nm/ps) are included
    /// so that the template is also usable as a pre-equilibrated starting frame.
    pub fn to_gro(&self) -> String {
        // OPC virtual-site coefficients: r_MW = r_OW + a*(r_HW1-r_OW) + a*(r_HW2-r_OW)
        const VS_A: f32 = 0.147803;

        let n_mols = self.o_posits.len();
        let n_atoms = n_mols * 4; // OW + HW1 + HW2 + MW

        let mut s = String::from("OPC water template — generated by Bio Files\n");
        let _ = writeln!(s, "{}", n_atoms);

        let mut atom_serial = 1usize;

        for i in 0..n_mols {
            let res = (i + 1) % 100_000;

            let ow = self.o_posits[i];
            let hw1 = self.h0_posits[i];
            let hw2 = self.h1_posits[i];
            let mw = ow + (hw1 - ow) * VS_A + (hw2 - ow) * VS_A;

            let v_ow = self.o_velocities[i];
            let v_hw1 = self.h0_velocities[i];
            let v_hw2 = self.h1_velocities[i];

            for (name, pos, vel) in [
                ("OW", ow, v_ow),
                ("HW1", hw1, v_hw1),
                ("HW2", hw2, v_hw2),
                ("MW", mw, Vec3::new_zero()), // virtual site — no independent velocity
            ] {
                let _ = writeln!(
                    s,
                    "{:>5}{:<5}{:>5}{:>5}{:>8.3}{:>8.3}{:>8.3}{:>8.4}{:>8.4}{:>8.4}",
                    res,
                    "SOL",
                    name,
                    atom_serial % 100_000,
                    pos.x / 10.0, // Å → nm
                    pos.y / 10.0,
                    pos.z / 10.0,
                    vel.x / 10.0, // Å/ps → nm/ps
                    vel.y / 10.0,
                    vel.z / 10.0,
                );
                atom_serial += 1;
            }
        }

        // Box vector line (Å → nm)
        let lo = self.bounds.0;
        let hi = self.bounds.1;
        let _ = writeln!(
            s,
            "{:>10.5}{:>10.5}{:>10.5}",
            (hi.x - lo.x).abs() / 10.0,
            (hi.y - lo.y).abs() / 10.0,
            (hi.z - lo.z).abs() / 10.0,
        );

        s
    }
}

/// Water model to use for explicit solvation via `gmx solvate`.
/// [solvate docs](https://manual.gromacs.org/current/onlinehelp/gmx-solvate.html#gmx-solvate)
///
/// Selecting a model causes the pipeline to:
/// 1. Include the model's atom-type and molecule-type sections in the topology.
/// 2. Run `gmx solvate` to fill the simulation box with water before `grompp`.
#[derive(Clone, Debug, PartialEq)]
pub enum WaterModel {
    /// OPC 4-point rigid water (recommended with Amber ff19SB).
    /// Requires GROMACS 2022+ which ships `opc.gro` in its data directory.
    Opc(WaterInitTemplate),
}

impl WaterModel {
    /// Pre-equilibrated water-box filename written by `save()` and passed to `gmx solvate -cs`.
    pub(in crate::gromacs) fn gro_filename(&self) -> &'static str {
        match self {
            Self::Opc(_) => SOLVENT_GRO_NAME,
        }
    }
}

/// Topology: Complete `SOL` molecule-type block for OPC 4-point water.
///
/// Geometry:
/// - d(O–H) = 0.08724 nm
/// - ∠H–O–H = 103.6° → d(H–H) = 0.13712 nm
/// - d(O–MW) = 0.01594 nm along the H–O–H bisector → virtual-site coefficients a = b = 0.147803
///
/// `gmx solvate -p` appends the molecule count; we only define the type here.
pub(in crate::gromacs) fn opc_sol_moleculetype() -> String {
    String::from(
        "[ moleculetype ]\n\
         ; molname  nrexcl\n\
         SOL        2\n\
         \n\
         [ atoms ]\n\
         ; nr  type    resnr  residue  atom  cgnr   charge     mass\n\
            1  OW_opc  1      SOL      OW    1       0.0000   15.99940\n\
            2  HW_opc  1      SOL      HW1   2       0.6791    1.00800\n\
            3  HW_opc  1      SOL      HW2   3       0.6791    1.00800\n\
            4  MW      1      SOL      MW    4      -1.3582    0.00000\n\
         \n\
         [ settles ]\n\
         ; OW  funct  dOH (nm)  dHH (nm)\n\
            1   1      0.08724   0.13712\n\
         \n\
         [ virtual_sites3 ]\n\
         ; MW  OW  HW1  HW2  funct  a         b\n\
            4   1    2    3   1     0.147803  0.147803\n\
         \n\
         [ exclusions ]\n\
         1  2  3  4\n\
         2  1  3  4\n\
         3  1  2  4\n\
         4  1  2  3\n\
         \n",
    )
}