dynamics 0.1.8

#![allow(clippy::excessive_precision)]

//! Code for initializing solvent molecules, including assigning quantity, initial positions, and
//! velocities. Set up to meet density, pressure, and or temperature targets. Not specific to the
//! solvent model used.
//!
//! This involves creating, saving and loading templates, and generating water molecules given a template,
//! sim box, and solute.

use std::{f32::consts::TAU, fs, io, path::Path, time::Instant};

use bincode::{Decode, Encode};
use bio_files::gromacs;
use lin_alg::{
    f32::{Mat3 as Mat3F32, Quaternion, Vec3},
    f64::{Quaternion as QuaternionF64, Vec3 as Vec3F64},
};
use rand::{Rng, distr::Uniform, rngs::ThreadRng};
use rand_distr::{Distribution, Normal};

use crate::{
    AtomDynamics, ComputationDevice, MdState, MolDynamics, NATIVE_TO_KCAL,
    barostat::SimBox,
    partial_charge_inference::{files::load_from_bytes_bincode, save},
    sa_surface,
    solvent::WaterMol,
    thermostat::{GAS_CONST_R, KB_A2_PS2_PER_K_PER_AMU},
};

// 0.997 g cm⁻³ is a good default density for biological pressures. We use this for initializing
// and maintaining the solvent density and molecule count.
const WATER_DENSITY: f32 = 0.997;

// g / mol (or AMU per molecule)
// This is similar to the Amber H and O masses we used summed, and could be explained
// by precision limits. We use it for generating atoms based on mass density.
const MASS_WATER: f32 = 18.015_28;

// Avogadro's constant. mol^-1.
const N_A: f32 = 6.022_140_76e23;

// This is ~0.0333 mol/ų
// Multiplying this by volume in Angstrom^3 gives us AMU g cm^-3 Å^3 mol^-1
const WATER_MOLS_PER_VOL: f32 = WATER_DENSITY * N_A / (MASS_WATER * 1.0e24);

// Don't generate solvent molecules that are too close to other atoms.
// Vdw contact distance between solvent molecules and organic molecules is roughly 3.5 Å.
// todo: Hmm. We could get lower, but there's some risk of an H atom being too close,
// todo and we're currently only measuring solvent o dist to atoms.
const MIN_NONWATER_DIST: f32 = 1.7;
const MIN_NONWATER_DIST_SQ: f32 = MIN_NONWATER_DIST * MIN_NONWATER_DIST;

// Direct O-O overlap check — prevents truly coincident molecules.
const MIN_WATER_O_O_DIST: f32 = 1.7;
const MIN_WATER_O_O_DIST_SQ: f32 = MIN_WATER_O_O_DIST * MIN_WATER_O_O_DIST;

// PBC-boundary exclusion distance.
// When a smaller box is filled from a larger template (e.g. 30 Å from a 60 Å template),
// molecules near opposite faces become PBC neighbours even though they were ~30 Å apart in
// the template and were never equilibrated at that short distance.  These pairs can slip
// through the 1.7 Å check with PBC distances of 2.0–2.8 Å, where LJ energy is 6–74 kcal/mol
// per pair — enough to push pressure into the tens-of-thousands-of-bar range.
// 2.8 Å is just below the first RDF peak (~2.82 Å).  We only apply this stricter threshold
// when PBC wrapping actually shortens the distance (i.e. `min_image_dist < direct_dist`),
// so interior template molecules at their natural 2.5–2.8 Å first-shell distances are
// accepted while un-equilibrated cross-boundary pairs are rejected.
const PBC_MIN_WATER_O_O_DIST: f32 = 2.8;
const PBC_MIN_WATER_O_O_DIST_SQ: f32 = PBC_MIN_WATER_O_O_DIST * PBC_MIN_WATER_O_O_DIST;

// Higher is more accurate, but slower. After hydrogen bond networks are settled, higher doensn't
// improve things. Note that we initialize from a pre-equilibrated template, so we shouldn't
// need many effects. This mainly deals with template tiling effects, and solvent-solute conflicts.
const NUM_EQUILIBRATION_STEPS: usize = 200;
// Like in our normal setup with constraint H, 0.002ps may be the safe upper bound.
// We seem to get better settling results with a low dt.
const DT_EQUILIBRATION: f32 = 0.0005;

// We use this externally, for example, when passing it to GROMACS in Molchanica.
pub const WATER_TEMPLATE_60A: &[u8] =
    include_bytes!("../../param_data/water_60A.water_init_template");

/// We store pre-equilibrated solvent molecules in a template, and use it to initialize solvent for a simulation.
/// This keeps the equilibration steps relatively low. Note that edge effects from tiling will require
/// equilibration, as well as adjusting a template for the runtime temperature target.
///
/// Struct-of-array layout. (Corresponding indices)
/// Public so it can be created by the application after a run.
///
/// 108 bytes/mol. Size on disk/mem: for a 60Å side len: ~780kb. (Hmm: We're getting a bit less)
/// 80Å/side: 1.20Mb.
#[derive(Encode, Decode)]
pub struct WaterInitTemplate {
    // velocity is o velocity, instead of 3 separate velocities
    o_posits: Vec<Vec3>,
    h0_posits: Vec<Vec3>,
    h1_posits: Vec<Vec3>,
    o_velocities: Vec<Vec3>,
    h0_velocities: Vec<Vec3>,
    h1_velocities: Vec<Vec3>,
    // todo: Cache these, or infer?
    /// This must correspond to the positions. Cached.
    bounds: (Vec3, Vec3),
}

impl WaterInitTemplate {
    /// Load a previously-saved template from a file path.
    pub fn load(path: &Path) -> io::Result<Self> {
        let bytes = fs::read(path)?;
        load_from_bytes_bincode(&bytes)
    }

    pub fn from_bytes(bytes: &[u8]) -> io::Result<Self> {
        load_from_bytes_bincode(bytes)
    }

    /// Construct from the current state, and save to file.
    /// Call this explicitly. (todo: Determine a formal or informal approach)
    pub fn create_and_save(
        water: &[WaterMol],
        bounds: (Vec3, Vec3),
        path: &Path,
    ) -> io::Result<()> {
        let n = water.len();

        let mut o_posits = Vec::with_capacity(n);
        let mut h0_posits = Vec::with_capacity(n);
        let mut h1_posits = Vec::with_capacity(n);

        let mut o_velocities = Vec::with_capacity(n);
        let mut h0_velocities = Vec::with_capacity(n);
        let mut h1_velocities = Vec::with_capacity(n);

        // let (mut min, mut max) = (Vec3::splat(f32::INFINITY), Vec3::splat(f32::NEG_INFINITY));

        // Sort solvent by position so that it iterates out from the center. This makes initialization
        // easier for cases where this template is larger than the target sim box.
        let water = {
            let ctr = (bounds.1 + bounds.0) / 2.;

            let mut w = water.to_vec();
            w.sort_by(|a, b| {
                let da = (a.o.posit - ctr).magnitude_squared();
                let db = (b.o.posit - ctr).magnitude_squared();
                da.total_cmp(&db)
            });

            w
        };

        for mol in water {
            o_posits.push(mol.o.posit);
            h0_posits.push(mol.h0.posit);
            h1_posits.push(mol.h1.posit);

            o_velocities.push(mol.o.vel);
            h0_velocities.push(mol.h0.vel);
            h1_velocities.push(mol.h1.vel);

            // min = min.min(mol.o.posit);
            // max = max.max(mol.o.posit);
        }

        let result = Self {
            o_posits,
            h0_posits,
            h1_posits,
            o_velocities,
            h0_velocities,
            h1_velocities,
            bounds,
        };

        save(path, &result)
    }

    // todo: Identical structs; we could consolidate.
    /// Note: `gmx solvate`` handles tiling, centering, and solute deconfliction; we can
    /// send it the raw template.
    pub fn to_gromacs(&self) -> gromacs::solvate::WaterInitTemplate {
        gromacs::solvate::WaterInitTemplate {
            o_posits: self.o_posits.clone(),
            h0_posits: self.h0_posits.clone(),
            h1_posits: self.h1_posits.clone(),
            o_velocities: self.o_velocities.clone(),
            h0_velocities: self.h0_velocities.clone(),
            h1_velocities: self.h1_velocities.clone(),
            bounds: self.bounds,
        }
    }
}

/// Determine the number of solvent molecules to add, based on box size and solute.
fn n_water_mols(cell: &SimBox, solute_atoms: &[AtomDynamics]) -> usize {
    let cell_volume = cell.volume();
    let mol_volume = sa_surface::vol_take_up_by_atoms(solute_atoms);
    let free_vol = cell_volume - mol_volume;

    let dims = format!(
        "{}:.2 x {:.2} x {:.2}",
        (cell.bounds_high.x - cell.bounds_low.x).abs(),
        (cell.bounds_high.y - cell.bounds_low.y).abs(),
        (cell.bounds_high.z - cell.bounds_low.z).abs()
    );

    println!(
        "Solvent-free vol: {:.2} Cell vol: {:.2} (ų / 1,000). Dims: {dims} Å",
        free_vol / 1_000.,
        cell_volume / 1_000.
    );

    // Estimate free volume & n_mols from it
    (WATER_MOLS_PER_VOL * free_vol).round() as usize
}

/// Create solvent molecules from a template, tiling it as many times as needed to fill the cell.
/// Works for any cell size: smaller than, equal to, or larger than the template.
///
/// The template is always centered on the cell. For cells smaller than the template only tile
/// (0,0,0) contributes; for larger cells neighboring tiles fill in the rest.
/// Water-solvent conflict detection uses min-image distances so molecules are never placed too
/// close to a PBC image of an already-placed molecule.
///
/// `template_override`: if provided, use this template instead of the built-in 60 Å one.
pub fn make_water_mols(
    cell: &SimBox,
    atoms: &[AtomDynamics],
    specify_num_water: Option<usize>,
    template_override: Option<&WaterInitTemplate>,
    // When true, skip the PBC-boundary proximity check (the 2.8 Å cross-boundary filter).
    // Only the hard-overlap 1.7 Å direct-distance check remains.  Use this when generating
    // a template at the correct equilibrium density: the ~88 boundary molecules that the PBC
    // filter would reject are acceptable starting points; the MD equilibration run will push
    // them to their natural first-shell distances.
    skip_pbc_filter: bool,
) -> Vec<WaterMol> {
    println!("Initializing solvent molecules...");
    let start = Instant::now();

    let default_template;

    let template: &WaterInitTemplate = match template_override {
        Some(t) => t,
        None => {
            default_template = load_from_bytes_bincode(WATER_TEMPLATE_60A).unwrap();
            &default_template
        }
    };

    let n_mols = specify_num_water.unwrap_or_else(|| n_water_mols(cell, atoms));
    let mut result = Vec::with_capacity(n_mols);

    if n_mols == 0 {
        println!("Complete in {} ms.", start.elapsed().as_millis());
        return result;
    }

    let atom_posits: Vec<_> = atoms.iter().map(|a| a.posit).collect();

    let template_size = template.bounds.1 - template.bounds.0;
    let template_ctr = (template.bounds.0 + template.bounds.1) / 2.;
    let cell_ctr = (cell.bounds_low + cell.bounds_high) / 2.;

    // Align tile (0,0,0) center to the cell center.
    let base_offset = cell_ctr - template_ctr;

    // Number of half-tiles needed to cover the cell in each direction (+1 for safety).
    let cell_size = cell.bounds_high - cell.bounds_low;
    let half_x = (cell_size.x / (2.0 * template_size.x)).ceil() as i32 + 1;
    let half_y = (cell_size.y / (2.0 * template_size.y)).ceil() as i32 + 1;
    let half_z = (cell_size.z / (2.0 * template_size.z)).ceil() as i32 + 1;

    let mut loops_used = 0;

    'tiles: for ix in -half_x..=half_x {
        for iy in -half_y..=half_y {
            for iz in -half_z..=half_z {
                let tile_offset = base_offset
                    + Vec3::new(
                        ix as f32 * template_size.x,
                        iy as f32 * template_size.y,
                        iz as f32 * template_size.z,
                    );

                'mol: for i in 0..template.o_posits.len() {
                    let o_posit = template.o_posits[i] + tile_offset;
                    let h0_posit = template.h0_posits[i] + tile_offset;
                    let h1_posit = template.h1_posits[i] + tile_offset;

                    loops_used += 1;

                    if !cell.contains(o_posit) {
                        continue;
                    }

                    // Conflict with solute atoms.
                    for &atom_p in &atom_posits {
                        if (atom_p - o_posit).magnitude_squared() < MIN_NONWATER_DIST_SQ {
                            continue 'mol;
                        }
                    }

                    // Conflict with already-placed solvent.
                    // Two-threshold check:
                    //   1. Direct distance < 1.7 Å: hard overlap regardless of PBC.
                    //   2. PBC-wrapped distance < 2.8 Å *and* wrapping shortened the distance:
                    //      these are cross-boundary pairs from the template that were never
                    //      equilibrated as PBC neighbours in this (smaller) cell.
                    //      Interior template molecules at their natural 2.5–2.8 Å first-shell
                    //      distances are not affected (min_image == direct for them).
                    for w in &result {
                        let diff = w.o.posit - o_posit;
                        let direct_sq = diff.magnitude_squared();
                        if direct_sq < MIN_WATER_O_O_DIST_SQ {
                            continue 'mol;
                        }
                        let min_image_sq = cell.min_image(diff).magnitude_squared();
                        // Always reject PBC hard overlaps (PBC distance < 1.7 Å) even when
                        // skip_pbc_filter is true, to prevent catastrophic initial forces.
                        if min_image_sq < MIN_WATER_O_O_DIST_SQ {
                            continue 'mol;
                        }
                        if !skip_pbc_filter {
                            if min_image_sq < PBC_MIN_WATER_O_O_DIST_SQ && min_image_sq < direct_sq
                            {
                                continue 'mol;
                            }
                        }
                    }

                    let mut mol = WaterMol::new(
                        Vec3::new_zero(),
                        Vec3::new_zero(),
                        Quaternion::new_identity(),
                    );
                    mol.o.posit = o_posit;
                    mol.h0.posit = h0_posit;
                    mol.h1.posit = h1_posit;
                    mol.o.vel = template.o_velocities[i];
                    mol.h0.vel = template.h0_velocities[i];
                    mol.h1.vel = template.h1_velocities[i];

                    result.push(mol);

                    if result.len() == n_mols {
                        break 'tiles;
                    }
                }
            }
        }
    }

    let elapsed = start.elapsed().as_millis();
    println!(
        "Added {} / {n_mols} solvent mols in {elapsed} ms. Used {loops_used} loops",
        result.len()
    );

    result
}

/// Pack copies of each custom solvent molecule into the simulation box, deconflicting with
/// already-placed atoms (e.g. the solute).  Returns one `MolDynamics` per copy, each with
/// `atom_posits` set to its chosen world-space positions.
///
/// Uses a greedy cubic-grid strategy with random-rotation search, identical in spirit to
/// `add_copies` in the molchanica layer.  Soft overlaps (> 1.4 Å apart) are accepted and
/// resolved later by the energy minimiser; hard overlaps are caught by `check_for_overlaps_oob`.
///
/// Only supports `SimBoxInit::Fixed` boxes — the caller is responsible for ensuring this.
pub(crate) fn pack_custom_solvent(
    bounds_low: Vec3,
    bounds_high: Vec3,
    existing_posits: &[Vec3F64], // declared positions of already-placed mols (e.g. solute)
    mols_solvent: &[(MolDynamics, usize)],
) -> Vec<MolDynamics> {
    // Below this squared distance between atoms we log a soft-overlap warning.
    // Energy minimisation resolves overlaps above the hard cutoff used in check_for_overlaps_oob.
    const MIN_ATOM_DIST_SQ: f64 = 1.4 * 1.4; // Ų
    // Keep every atom at least this far from each box face.
    const WALL_MARGIN: f64 = 0.6; // Å  (slightly > check_for_overlaps_oob's 0.5 Å limit)
    const MAX_ROT_ATTEMPTS: usize = 200;

    let mut rng = rand::rng();

    let lo = Vec3F64::new(
        bounds_low.x as f64,
        bounds_low.y as f64,
        bounds_low.z as f64,
    );
    let hi = Vec3F64::new(
        bounds_high.x as f64,
        bounds_high.y as f64,
        bounds_high.z as f64,
    );
    let box_size = hi - lo;
    let box_ctr = (lo + hi) * 0.5;

    // Grows as copies are committed; starts with the solute atom positions.
    let mut placed_posits: Vec<Vec3F64> = existing_posits.to_vec();

    let mut result: Vec<MolDynamics> = Vec::new();

    for (mol, count) in mols_solvent {
        let count = *count;
        if count == 0 {
            continue;
        }

        // Template positions in world space; prefer atom_posits override.
        let template_world: Vec<Vec3F64> = if let Some(ap) = &mol.atom_posits {
            ap.clone()
        } else {
            mol.atoms.iter().map(|a| a.posit).collect()
        };

        let n_atoms = template_world.len();
        if n_atoms == 0 {
            continue;
        }

        // Centroid and centroid-relative locals.
        let centroid = template_world
            .iter()
            .fold(Vec3F64::new(0., 0., 0.), |s, &p| s + p)
            * (1.0 / n_atoms as f64);
        let local: Vec<Vec3F64> = template_world.iter().map(|&p| p - centroid).collect();
        let bounding_r: f64 = local.iter().map(|p| p.magnitude()).fold(0.0_f64, f64::max);

        // Spatial early-reject radius: only check placed atoms within this of a candidate centroid.
        let search_sq = (bounding_r * 2.0 + 2.0).powi(2);

        // Every copy's centroid must be ≥ bounding_r + WALL_MARGIN from each wall face so that
        // even the molecule's furthest atom stays within the per-atom wall margin.
        let safe_margin = bounding_r + WALL_MARGIN;
        let inner_lo = lo + Vec3F64::new(safe_margin, safe_margin, safe_margin);
        let inner_hi = hi - Vec3F64::new(safe_margin, safe_margin, safe_margin);

        if inner_lo.x >= inner_hi.x || inner_lo.y >= inner_hi.y || inner_lo.z >= inner_hi.z {
            eprintln!(
                "pack_custom_solvent: box too small for molecule \
                 (bounding_r={:.1} Å, need >{:.1} Å per side); skipping {} copies.",
                bounding_r,
                2.0 * safe_margin,
                count
            );
            continue;
        }

        let inner_size = inner_hi - inner_lo;
        // Scale n up from ∛count so that valid grid cells (those within inner_lo..inner_hi)
        // number at least `count`.  Use the tightest dimension as the conservative factor.
        let naive_n = (count as f64).cbrt().ceil() as usize;
        let scale = (box_size.x / inner_size.x)
            .max(box_size.y / inner_size.y)
            .max(box_size.z / inner_size.z);
        let n = ((naive_n as f64 * scale).ceil() as usize).max(3);
        let (sx, sy, sz) = (
            box_size.x / n as f64,
            box_size.y / n as f64,
            box_size.z / n as f64,
        );
        // Box-centred half-widths for the per-atom wall check.
        let (hx, hy, hz) = (
            box_size.x * 0.5 - WALL_MARGIN,
            box_size.y * 0.5 - WALL_MARGIN,
            box_size.z * 0.5 - WALL_MARGIN,
        );

        // Grid cell centres restricted to the safe inner region.
        let mut grid: Vec<Vec3F64> = (0..n)
            .flat_map(|ix| {
                (0..n).flat_map(move |iy| {
                    (0..n).map(move |iz| {
                        Vec3F64::new(
                            lo.x + (ix as f64 + 0.5) * sx,
                            lo.y + (iy as f64 + 0.5) * sy,
                            lo.z + (iz as f64 + 0.5) * sz,
                        )
                    })
                })
            })
            .filter(|c| {
                c.x >= inner_lo.x
                    && c.x <= inner_hi.x
                    && c.y >= inner_lo.y
                    && c.y <= inner_hi.y
                    && c.z >= inner_lo.z
                    && c.z <= inner_hi.z
            })
            .collect();

        for copy_i in 0..count {
            // Greedy: pick the cell whose centroid is furthest from all placed atoms.
            let best_cell_idx = if placed_posits.is_empty() {
                0
            } else {
                grid.iter()
                    .enumerate()
                    .map(|(i, &cell_ctr)| {
                        let min_dsq = placed_posits
                            .iter()
                            .map(|&p| (cell_ctr - p).magnitude_squared())
                            .fold(f64::MAX, f64::min);
                        (i, min_dsq)
                    })
                    .max_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal))
                    .map(|(i, _)| i)
                    .unwrap_or(0)
            };

            let world_ctr = grid.remove(best_cell_idx);

            let mut best_min_sq = f64::NEG_INFINITY;
            let mut best_posits: Vec<Vec3F64> = Vec::new();

            for _ in 0..MAX_ROT_ATTEMPTS {
                let (w, x, y, z): (f64, f64, f64, f64) =
                    (rng.random(), rng.random(), rng.random(), rng.random());
                let rot = QuaternionF64::new(w, x, y, z).to_normalized();

                let new_posits: Vec<Vec3F64> = local
                    .iter()
                    .map(|&l| rot.rotate_vec(l) + world_ctr)
                    .collect();

                // Wall check in box-centred coordinates.
                if !new_posits.iter().all(|p| {
                    let dp = *p - box_ctr;
                    dp.x.abs() <= hx && dp.y.abs() <= hy && dp.z.abs() <= hz
                }) {
                    continue;
                }

                // Atom-level clash check against placed atoms.
                let mut min_sq = f64::MAX;
                'check: for &np in &new_posits {
                    for &pp in &placed_posits {
                        if (pp - world_ctr).magnitude_squared() > search_sq {
                            continue;
                        }
                        let dsq = (np - pp).magnitude_squared();
                        if dsq < min_sq {
                            min_sq = dsq;
                            if min_sq < MIN_ATOM_DIST_SQ {
                                break 'check;
                            }
                        }
                    }
                }

                if min_sq > best_min_sq {
                    best_min_sq = min_sq;
                    best_posits = new_posits;
                }
                if best_min_sq >= MIN_ATOM_DIST_SQ {
                    break; // Clean placement found.
                }
            }

            // Fallback if every rotation attempt failed the wall check.
            if best_posits.is_empty() {
                best_posits = local.iter().map(|&l| l + world_ctr).collect();
            }

            if best_min_sq < MIN_ATOM_DIST_SQ {
                eprintln!(
                    "pack_custom_solvent: copy {copy_i}: best min atom dist {:.2} Å — \
                     placing with soft overlap (energy minimiser will resolve).",
                    best_min_sq.max(0.0).sqrt()
                );
            }

            placed_posits.extend_from_slice(&best_posits);

            let mut mol_copy = mol.clone();
            mol_copy.atom_posits = Some(best_posits);
            result.push(mol_copy);

            if grid.is_empty() && copy_i + 1 < count {
                eprintln!(
                    "pack_custom_solvent: grid cells exhausted after {} / {} copies; \
                     box may be too small for this many solvent molecules.",
                    copy_i + 1,
                    count
                );
                break;
            }
        }
    }

    result
}

#[allow(unused)]
/// Creates a regular lattice of water molecules. We use this as the first part of creating
/// a solvent template. Use this,  run a sim with thermostat and barostat, then store the result
/// in a `WaterInitTemplate`. We can save and load this to disk as binary, or in `.gro` format.
///
/// Generate solvent molecules to meet a temperature target, using standard density assumptions.
/// We deconflict with (solute) atoms in the simulation, and base the number of molecules to add
/// on the free space, not the total cell volume.
///
/// Process:
/// - Compute the number of molecules to add
/// - Add them on a regular grid with random orientations, and velocities in a random distribution
///   that matches the target temperature. Move molecules to the edge that are too close to
///   solute atoms.
///
/// Note: If we're able to place most, but not all waters, the barostat should adjust the sim box size
/// to account for the lower-than-specific pressure.
pub fn make_water_mols_grid(
    cell: &SimBox,
    temperature_tgt: f32,
    zero_com_drift: bool,
) -> Vec<WaterMol> {
    println!("Initializing a solvent grid, as part of template preparation...");
    // Initialize an RNG for orientations.
    let mut rng = rand::rng();
    let distro = Uniform::<f32>::new(0.0, 1.0).unwrap();

    let n_mols = n_water_mols(cell, &[]);

    let mut result: Vec<WaterMol> = Vec::with_capacity(n_mols);

    // Initialize the correct number of solvent molecules on a uniform grid. We ignore the solute for
    let lx = cell.bounds_high.x - cell.bounds_low.x;
    let ly = cell.bounds_high.y - cell.bounds_low.y;
    let lz = cell.bounds_high.z - cell.bounds_low.z;

    let base = (n_mols as f32).cbrt().round().max(1.0) as usize;
    let n_x = base;
    let n_y = base;
    let n_z = n_mols.div_ceil(n_x * n_y);

    let spacing_x = lx / n_x as f32;
    let spacing_y = ly / n_y as f32;
    let spacing_z = lz / n_z as f32;

    // Prevents unbounded looping. A higher value means we're more likely to succed,
    // but the run time could be higher.
    let fault_ratio = 3;

    let mut num_added = 0;
    let mut loops_used = 0;

    'outer: for i in 0..n_mols * fault_ratio {
        let a = i % n_x;
        let b = (i / n_x) % n_y;
        let c = (i / (n_x * n_y)) % n_z;

        let posit = Vec3::new(
            cell.bounds_low.x + (a as f32 + 0.5) * spacing_x,
            cell.bounds_low.y + (b as f32 + 0.5) * spacing_y,
            cell.bounds_low.z + (c as f32 + 0.5) * spacing_z,
        );

        // Check for an overlap with existing solvent molecules.
        for w in &result {
            let dist_sq = (w.o.posit - posit).magnitude_squared();
            if dist_sq < MIN_WATER_O_O_DIST_SQ {
                loops_used += 1;
                continue 'outer;
            }
        }

        result.push(WaterMol::new(
            posit,
            Vec3::new_zero(),
            random_quaternion(&mut rng, distro),
        ));
        num_added += 1;

        if num_added == n_mols {
            break;
        }
        loops_used += 1;
    }

    // Set velocities consistent with the temperature target.
    init_velocities(&mut result, temperature_tgt, zero_com_drift, &mut rng);

    println!(
        "Added {} / {n_mols} solvent mols. Used {loops_used} loops",
        result.len()
    );
    result
}

/// We use this as part of our water template generation.
///
/// Note: This sets a reasonable default, but our thermostat, applied notably during
/// our initial solvent simulation, determines the actual temperature set at proper sim init.
/// Note: We've deprecated this in favor of velocities pre-initialized in the template.
fn init_velocities(
    mols: &mut [WaterMol],
    t_target: f32,
    zero_com_drift: bool,
    rng: &mut ThreadRng,
) {
    let kT = KB_A2_PS2_PER_K_PER_AMU * t_target;

    for m in mols.iter_mut() {
        // COM and relative positions
        let (r_com, m_tot) = {
            let mut r = Vec3::new_zero();
            let mut m_tot = 0.0;
            for a in [&m.o, &m.h0, &m.h1] {
                r += a.posit * a.mass;
                m_tot += a.mass;
            }
            (r / m_tot, m_tot)
        };

        let r_0 = m.o.posit - r_com;
        let r_h0 = m.h0.posit - r_com;
        let r_h1 = m.h1.posit - r_com;

        // Sample COM velocity
        let sigma_v = (kT / m_tot).sqrt();
        let n = Normal::new(0.0, sigma_v).unwrap();
        let v_com = Vec3::new(n.sample(rng), n.sample(rng), n.sample(rng));

        // Inertia tensor about COM (world frame)
        // Build as arrays (your code)
        let inertia = |r: Vec3, mass: f32| {
            let r2 = r.dot(r);
            [
                [
                    mass * (r2 - r.x * r.x),
                    -mass * r.x * r.y,
                    -mass * r.x * r.z,
                ],
                [
                    -mass * r.y * r.x,
                    mass * (r2 - r.y * r.y),
                    -mass * r.y * r.z,
                ],
                [
                    -mass * r.z * r.x,
                    -mass * r.z * r.y,
                    mass * (r2 - r.z * r.z),
                ],
            ]
        };
        let mut I_arr = inertia(r_0, m.o.mass);
        let add_I = |I: &mut [[f32; 3]; 3], J: [[f32; 3]; 3]| {
            for i in 0..3 {
                for j in 0..3 {
                    I[i][j] += J[i][j];
                }
            }
        };
        add_I(&mut I_arr, inertia(r_h0, m.h0.mass));
        add_I(&mut I_arr, inertia(r_h1, m.h1.mass));

        let I = Mat3F32::from_arr(I_arr);

        // Diagonalize and solve with the Mat3 methods
        let (eigvecs, eigvals) = I.eigen_vecs_vals();
        let L_principal = Vec3::new(
            Normal::new(0.0, (kT * eigvals.x.max(0.0)).sqrt())
                .unwrap()
                .sample(rng),
            Normal::new(0.0, (kT * eigvals.y.max(0.0)).sqrt())
                .unwrap()
                .sample(rng),
            Normal::new(0.0, (kT * eigvals.z.max(0.0)).sqrt())
                .unwrap()
                .sample(rng),
        );
        let L_world = eigvecs * L_principal; // assumes Mat3 * Vec3 is implemented
        let omega = I.solve_system(L_world); // ω = I^{-1} L

        // Set atomic velocities
        m.o.vel = v_com + omega.cross(r_0);
        m.h0.vel = v_com + omega.cross(r_h0);
        m.h1.vel = v_com + omega.cross(r_h1);
    }

    if zero_com_drift {
        // Remove global COM drift
        remove_com_velocity(mols);
    }

    let (ke_raw, dof) = _kinetic_energy_and_dof(mols, zero_com_drift);

    // current T = 2 KE / (dof * R)
    let temperature_meas = (2.0 * ke_raw) / (dof as f32 * GAS_CONST_R as f32);
    let lambda = (t_target / temperature_meas).sqrt();

    for a in atoms_mut(mols) {
        if a.mass > 0.0 {
            a.vel *= lambda;
        }
    }
}

/// Calculate kinetic energy in kcal/mol, and DOF for solvent only.
/// Water is rigid, so 3 DOF per molecule.
fn _kinetic_energy_and_dof(mols: &[WaterMol], zero_com_drift: bool) -> (f32, usize) {
    let mut ke = 0.;
    for w in mols {
        ke += (w.o.mass * w.o.vel.magnitude_squared()) as f64;
        ke += (w.h0.mass * w.h0.vel.magnitude_squared()) as f64;
        ke += (w.h1.mass * w.h1.vel.magnitude_squared()) as f64;
    }

    let mut dof = mols.len() * 3;

    if zero_com_drift {
        dof = dof.saturating_sub(3);
    }

    // Add in the 0.5 factor, and convert from amu • (Å/ps)² to kcal/mol.
    (ke as f32 * 0.5 * NATIVE_TO_KCAL, dof)
}

fn atoms_mut(mols: &mut [WaterMol]) -> impl Iterator<Item = &mut AtomDynamics> {
    mols.iter_mut()
        .flat_map(|m| [&mut m.o, &mut m.h0, &mut m.h1].into_iter())
}

#[allow(unused)]
/// Removes center-of-mass drift. Use in template generation
fn remove_com_velocity(mols: &mut [WaterMol]) {
    let mut p = Vec3::new_zero();
    let mut m_tot = 0.0;
    for a in atoms_mut(mols) {
        p += a.vel * a.mass;
        m_tot += a.mass;
    }

    let v_com = p / m_tot;
    for a in atoms_mut(mols) {
        a.vel -= v_com;
    }
}

#[allow(unused)]
/// Used in template generation
// todo: It might be nice to have this in lin_alg, although I don't want to add the rand
// todo dependency to it.
fn random_quaternion(rng: &mut ThreadRng, distro: Uniform<f32>) -> Quaternion {
    let (u1, u2, u3) = (rng.sample(distro), rng.sample(distro), rng.sample(distro));
    let sqrt1_minus_u1 = (1.0 - u1).sqrt();
    let sqrt_u1 = u1.sqrt();
    let (theta1, theta2) = (TAU * u2, TAU * u3);

    Quaternion::new(
        sqrt1_minus_u1 * theta1.sin(),
        sqrt1_minus_u1 * theta1.cos(),
        sqrt_u1 * theta2.sin(),
        sqrt_u1 * theta2.cos(),
    )
    .to_normalized()
}

impl MdState {
    /// Use this to help initialize solvent molecules to realistic geometry of hydrogen bond networks,
    /// prior to the first proper simulation step. Runs MD on solvent only.
    /// Make sure to only run this after state is properly initialized, e.g. towards the end
    /// of init; not immediately after populating waters.
    ///
    /// This will result in an immediate energy bump as solvent positions settle from their grid
    /// into position. As they settle, the thermostat will bring the velocities down to set
    /// the target temp. This sim should run long enough to the solvent is stable by the time
    /// the main sim starts.
    pub fn md_on_water_only(&mut self, dev: &ComputationDevice) {
        println!("Initializing solvent H bond networks...");
        let start = Instant::now();

        // This disables things like snapshot saving, and certain prints.
        self.solvent_only_sim_at_init = true;

        // Mark all non-solvent atoms as static; keep track of their original state here.
        let mut static_state = Vec::with_capacity(self.atoms.len());
        for a in &mut self.atoms {
            static_state.push(a.static_);
            a.static_ = true;
        }

        for _ in 0..NUM_EQUILIBRATION_STEPS {
            self.step(dev, DT_EQUILIBRATION, None);
        }

        // Restore the original static state.
        for (i, a) in self.atoms.iter_mut().enumerate() {
            a.static_ = static_state[i];
        }

        self.solvent_only_sim_at_init = false;
        self.step_count = 0; // Reset.

        let elapsed = start.elapsed().as_millis();
        println!("Water H bond networks complete in {elapsed} ms");
    }
}