bio_files 0.5.0

//! GROMACS topology (`.top`) generation from Amber-format (Or dyanmics library) parameters.
//!
//! [Example .top file](https://manual.gromacs.org/2026.1/reference-manual/file-formats.html#top)
//! [User Guide: Force Fields](https://manual.gromacs.org/current/user-guide/force-fields.html)
//!
//! Produces a fully self-contained topology that does not depend on any external
//! GROMACS force-field include files. All bonded and non-bonded parameters are
//! inlined from the supplied [`ForceFieldParams`].
//!
//! ## Unit conversions (Amber → GROMACS)
//!
//! | Quantity       | Amber unit        | GROMACS unit        | Factor          |
//! |----------------|-------------------|---------------------|-----------------|
//! | Length         | Å                 | nm                  | ÷ 10            |
//! | Energy         | kcal/mol          | kJ/mol              | × 4.184         |
//! | Bond *k*       | kcal/mol/Ų       | kJ/mol/nm²          | × 418.4         |
//! | Angle *k*      | kcal/mol/rad²     | kJ/mol/rad²         | × 4.184         |
//! | Dihedral *V*   | kcal/mol          | kJ/mol              | × 4.184         |
//! | LJ sigma       | Å                 | nm                  | ÷ 10            |
//! | LJ epsilon     | kcal/mol          | kJ/mol              | × 4.184         |

use std::collections::{HashMap, HashSet};

use crate::{
    AtomGeneric, BondGeneric,
    gromacs::{solvate, solvate::WaterModel},
    md_params::ForceFieldParams,
};

const KCAL_TO_KJ: f32 = 4.184;
const ANG_TO_NM: f32 = 0.1;
/// kcal/mol/Ų → kJ/mol/nm²
const BOND_K_FACTOR: f32 = KCAL_TO_KJ * 100.0;

// ---------------------------------------------------------------------------
// Public entry point
// ---------------------------------------------------------------------------

/// One molecule type entry — atoms, bonds, a name, and an optional copy count.
pub struct MoleculeTopology<'a> {
    pub name: &'a str,
    pub atoms: &'a [AtomGeneric],
    pub bonds: &'a [BondGeneric],
    /// Per-molecule force-field parameters (e.g. GAFF2 for a ligand).
    /// Falls back to `ff_global` when a term is missing.
    pub ff_mol: Option<&'a ForceFieldParams>,
    /// Number of copies listed in `[ molecules ]`.
    pub count: usize,
}

/// Generate a complete, self-contained GROMACS `.top` string.
///
/// `ff_global` is the system-wide parameter set (e.g. ff19SB / GAFF2 merged), used
/// as a fall-back when a term is absent from the per-molecule `ff_mol`.
pub fn make_top(
    molecules: &[MoleculeTopology<'_>],
    ff_global: Option<&ForceFieldParams>,
    water_model: Option<&WaterModel>,
) -> String {
    let mut s = String::from("; GROMACS topology generated by Bio Files\n\n");

    // --- [ defaults ] -------------------------------------------------------
    // nbfunc 1  = Lennard-Jones
    // comb-rule 2 = Lorentz–Berthelot (arithmetic sigma, geometric epsilon)
    // fudgeLJ / fudgeQQ: Amber 1-4 scaling factors
    s.push_str("[ defaults ]\n");
    s.push_str("; nbfunc  comb-rule  gen-pairs  fudgeLJ  fudgeQQ\n");
    s.push_str("  1       2          yes        0.5      0.8333\n\n");

    // --- [ atomtypes ] -------------------------------------------------------
    // Collect the union of all FF types referenced across all molecules.
    let mut all_types: HashSet<String> = HashSet::new();
    for mol in molecules {
        for atom in mol.atoms {
            if let Some(t) = &atom.force_field_type {
                all_types.insert(t.clone());
            }
        }
    }

    s.push_str("[ atomtypes ]\n");
    s.push_str("; name  at.num   mass      charge  ptype  sigma (nm)      epsilon (kJ/mol)\n");
    for ff_type in &all_types {
        // Try per-molecule params first, then the global fallback — same two-source
        // chain used by lookup_bond/angle/dihedral. This matters because ff_mol often
        // carries only bonded terms, while mass/LJ live in ff_global (e.g. GAFF2 table).
        let mass_val = molecules
            .iter()
            .find_map(|m| m.ff_mol.and_then(|p| p.mass.get(ff_type.as_str())))
            .or_else(|| ff_global.and_then(|p| p.mass.get(ff_type.as_str())))
            .map(|m| m.mass)
            .unwrap_or(12.011);

        let (sigma_nm, eps_kj) = molecules
            .iter()
            .find_map(|m| m.ff_mol.and_then(|p| p.lennard_jones.get(ff_type.as_str())))
            .or_else(|| ff_global.and_then(|p| p.lennard_jones.get(ff_type.as_str())))
            .map(|lj| (lj.sigma * ANG_TO_NM, lj.eps * KCAL_TO_KJ))
            .unwrap_or((0.3, 0.5));

        let at_num = atomic_number_from_mass(mass_val);

        s.push_str(&format!(
            "  {:<6}  {:>3}    {:>8.4}   0.000  A  {:>14.8e}  {:>14.8e}\n",
            ff_type, at_num, mass_val, sigma_nm, eps_kj,
        ));
    }

    // Water model atom types
    if let Some(WaterModel::Opc(_)) = water_model {
        s.push_str(&opc_atomtypes());
    }
    s.push('\n');

    // --- Per-molecule blocks -------------------------------------------------
    for mol in molecules {
        write_molecule_block(&mut s, mol, ff_global);
    }

    // Water model molecule type (SOL) — must appear before [system]/[molecules]
    // so grompp can resolve it after gmx solvate appends "SOL N" to [molecules].
    if let Some(WaterModel::Opc(_)) = water_model {
        s.push_str(&solvate::opc_sol_moleculetype());
    }

    // --- [ system ] ----------------------------------------------------------
    s.push_str("[ system ]\n; Name\nMolchanica MD\n\n");

    // --- [ molecules ] -------------------------------------------------------
    // Solute only — gmx solvate appends "SOL N" automatically when run with -p.
    s.push_str("[ molecules ]\n; Compound     #mols\n");
    for mol in molecules {
        s.push_str(&format!("{:<14}  {}\n", mol.name, mol.count));
    }

    s
}

// ---------------------------------------------------------------------------
// OPC water topology helpers
// ---------------------------------------------------------------------------

/// Atom-type lines for OPC water to append to the `[atomtypes]` section.
///
/// Parameters from GROMACS `opc.itp` (GROMACS 2022+):
/// - OW_opc: oxygen, LJ site only (no charge on O in OPC)
/// - HW_opc: hydrogen, no LJ
/// - MW:     massless virtual charge site (4th point)
fn opc_atomtypes() -> String {
    // sigma in nm, epsilon in kJ/mol (from Amber frcmod.opc converted to GROMACS units)
    // OW_opc: R* = 1.777167 Å → sigma = R* * 2/2^(1/6) = 0.316655 nm; eps = 0.2128 kcal/mol = 0.8907 kJ/mol
    String::from(concat!(
        "  OW_opc    8   15.99940   0.0000  A   3.16655e-01  8.90699e-01\n",
        "  HW_opc    1    1.00800   0.0000  A   0.00000e+00  0.00000e+00\n",
        "  MW        0    0.00000   0.0000  V   0.00000e+00  0.00000e+00\n",
    ))
}

// ---------------------------------------------------------------------------
// Internal helpers
// ---------------------------------------------------------------------------

// ---------------------------------------------------------------------------
// Parameter look-ups with symmetric key fallback
// ---------------------------------------------------------------------------

pub(in crate::gromacs) fn lookup_bond(
    i: &str,
    j: &str,
    ff: Option<&ForceFieldParams>,
    ff_global: Option<&ForceFieldParams>,
) -> Option<(f32, f32)> {
    let key_fwd = (i.to_string(), j.to_string());
    let key_rev = (j.to_string(), i.to_string());

    for src in [ff, ff_global].into_iter().flatten() {
        if let Some(p) = src.bond.get(&key_fwd).or_else(|| src.bond.get(&key_rev)) {
            return Some((p.r_0, p.k_b));
        }
    }
    None
}

pub(in crate::gromacs) fn lookup_angle(
    i: &str,
    j: &str,
    k: &str,
    ff: Option<&ForceFieldParams>,
    ff_global: Option<&ForceFieldParams>,
) -> Option<(f32, f32)> {
    let key_fwd = (i.to_string(), j.to_string(), k.to_string());
    let key_rev = (k.to_string(), j.to_string(), i.to_string());

    for src in [ff, ff_global].into_iter().flatten() {
        if let Some(p) = src.angle.get(&key_fwd).or_else(|| src.angle.get(&key_rev)) {
            return Some((p.theta_0, p.k));
        }
    }
    None
}

pub(in crate::gromacs) fn lookup_dihedral<'a>(
    i: &str,
    j: &str,
    k: &str,
    l: &str,
    ff: Option<&'a ForceFieldParams>,
    ff_global: Option<&'a ForceFieldParams>,
) -> Option<&'a Vec<crate::md_params::DihedralParams>> {
    let key_fwd = (i.to_string(), j.to_string(), k.to_string(), l.to_string());
    let key_rev = (l.to_string(), k.to_string(), j.to_string(), i.to_string());

    // Wildcard keys used by Amber ("X" placeholders)
    let key_wc_fwd = (
        "X".to_string(),
        j.to_string(),
        k.to_string(),
        "X".to_string(),
    );
    let key_wc_rev = (
        "X".to_string(),
        k.to_string(),
        j.to_string(),
        "X".to_string(),
    );

    for src in [ff, ff_global].into_iter().flatten() {
        for key in [&key_fwd, &key_rev, &key_wc_fwd, &key_wc_rev] {
            if let Some(p) = src.dihedral.get(key) {
                return Some(p);
            }
        }
    }
    None
}

// ---------------------------------------------------------------------------
// Topology traversal helpers
// ---------------------------------------------------------------------------

/// Build an adjacency list (1-based index → neighbours) from bonds.
fn adjacency(sn_to_idx: &HashMap<u32, usize>, bonds: &[BondGeneric]) -> HashMap<usize, Vec<usize>> {
    let mut adj: HashMap<usize, Vec<usize>> = HashMap::new();
    for bond in bonds {
        let (Some(&ai), Some(&aj)) = (
            sn_to_idx.get(&bond.atom_0_sn),
            sn_to_idx.get(&bond.atom_1_sn),
        ) else {
            continue;
        };
        adj.entry(ai).or_default().push(aj);
        adj.entry(aj).or_default().push(ai);
    }
    adj
}

/// Enumerate all 1-2-3 angle triples (each unique, centre atom = aj).
pub(in crate::gromacs) fn enumerate_angles(
    sn_to_idx: &HashMap<u32, usize>,
    bonds: &[BondGeneric],
) -> Vec<(usize, usize, usize)> {
    let adj = adjacency(sn_to_idx, bonds);
    let mut triples = Vec::new();
    let mut seen: HashSet<(usize, usize, usize)> = HashSet::new();

    for (&aj, neighbours) in &adj {
        let n = neighbours.len();
        for a in 0..n {
            for b in (a + 1)..n {
                let ai = neighbours[a];
                let ak = neighbours[b];
                let key = if ai < ak { (ai, aj, ak) } else { (ak, aj, ai) };
                if seen.insert(key) {
                    triples.push(key);
                }
            }
        }
    }
    triples
}

/// Enumerate all hub-and-spoke improper quads: (sat0, sat1, hub, sat2).
/// Hub is at index 2 (position 3) — Amber/GAFF convention.
/// Only atoms with ≥3 neighbours are candidates.
fn enumerate_impropers(
    sn_to_idx: &HashMap<u32, usize>,
    bonds: &[BondGeneric],
) -> Vec<(usize, usize, usize, usize)> {
    let adj = adjacency(sn_to_idx, bonds);
    let mut quads = Vec::new();
    let mut seen: HashSet<(usize, usize, usize, usize)> = HashSet::new();

    for (&hub, neighbors) in &adj {
        let n = neighbors.len();
        if n < 3 {
            continue;
        }
        for a in 0..n - 2 {
            for b in a + 1..n - 1 {
                for d in b + 1..n {
                    let (sat0, sat1, sat2) = (neighbors[a], neighbors[b], neighbors[d]);
                    let key = (sat0, sat1, hub, sat2);
                    if seen.insert(key) {
                        quads.push(key);
                    }
                }
            }
        }
    }
    quads
}

/// Look up an improper dihedral term.  Satellites are sorted alphabetically
/// before forming the key, matching Amber's wildcard-lookup convention.
fn lookup_improper<'a>(
    sat0: &str,
    sat1: &str,
    hub: &str,
    sat2: &str,
    ff: Option<&'a crate::md_params::ForceFieldParams>,
    ff_global: Option<&'a crate::md_params::ForceFieldParams>,
) -> Option<&'a Vec<crate::md_params::DihedralParams>> {
    let mut sats = [sat0, sat1, sat2];
    sats.sort_unstable();
    let key = (
        sats[0].to_string(),
        sats[1].to_string(),
        hub.to_string(),
        sats[2].to_string(),
    );
    for src in [ff, ff_global].into_iter().flatten() {
        if let Some(p) = src.get_dihedral(&key, false, true) {
            return Some(p);
        }
    }
    None
}

/// Enumerate all 1-2-3-4 proper dihedral quads.
fn enumerate_dihedrals(
    sn_to_idx: &HashMap<u32, usize>,
    bonds: &[BondGeneric],
) -> Vec<(usize, usize, usize, usize)> {
    let adj = adjacency(sn_to_idx, bonds);
    let mut quads = Vec::new();
    let mut seen: HashSet<(usize, usize, usize, usize)> = HashSet::new();

    for bond in bonds {
        let (Some(&aj), Some(&ak)) = (
            sn_to_idx.get(&bond.atom_0_sn),
            sn_to_idx.get(&bond.atom_1_sn),
        ) else {
            continue;
        };
        for &ai in adj.get(&aj).into_iter().flatten() {
            if ai == ak {
                continue;
            }
            for &al in adj.get(&ak).into_iter().flatten() {
                if al == aj || al == ai {
                    continue;
                }
                let key = if (ai, aj) < (al, ak) {
                    (ai, aj, ak, al)
                } else {
                    (al, ak, aj, ai)
                };
                if seen.insert(key) {
                    quads.push(key);
                }
            }
        }
    }
    quads
}

// ---------------------------------------------------------------------------
// Misc helpers
// ---------------------------------------------------------------------------

/// Very coarse atomic-number estimate from atomic mass — sufficient for the
/// informational `at.num` field in `[ atomtypes ]`, which GROMACS uses only for
/// visualisation and does not affect energetics.
pub fn atomic_number_from_mass(mass: f32) -> u8 {
    match mass as u32 {
        1 => 1,        // H
        4 => 2,        // He
        7 => 3,        // Li
        12 => 6,       // C
        14 => 7,       // N
        16 => 8,       // O
        19 => 9,       // F
        23 => 11,      // Na
        24 => 12,      // Mg
        31 => 15,      // P
        32 => 16,      // S
        35 | 36 => 17, // Cl
        39 => 19,      // K
        40 => 20,      // Ca
        56 => 26,      // Fe
        63 | 64 => 29, // Cu
        65 => 30,      // Zn
        _ => 6,        // default to C
    }
}

fn write_molecule_block(
    s: &mut String,
    mol: &MoleculeTopology<'_>,
    ff_global: Option<&ForceFieldParams>,
) {
    let ff = mol.ff_mol.or(ff_global);

    // Build serial-number → 1-based index map.
    let sn_to_idx: HashMap<u32, usize> = mol
        .atoms
        .iter()
        .enumerate()
        .map(|(i, a)| (a.serial_number, i + 1))
        .collect();

    // [ moleculetype ]
    s.push_str("[ moleculetype ]\n; molname  nrexcl\n");
    s.push_str(&format!("  {}        3\n\n", mol.name));

    // [ atoms ]
    s.push_str("[ atoms ]\n");
    s.push_str("; nr   type    resnr  residue  atom    cgnr   charge     mass\n");

    for (i, atom) in mol.atoms.iter().enumerate() {
        let nr = i + 1;
        let ff_type = atom.force_field_type.as_deref().unwrap_or("X");
        let charge = atom.partial_charge.unwrap_or(0.0);

        let mass = mol
            .ff_mol
            .and_then(|p| p.mass.get(ff_type))
            .or_else(|| ff_global.and_then(|p| p.mass.get(ff_type)))
            .map(|m| m.mass)
            .unwrap_or(12.011);

        let atom_name = if atom.hetero {
            // Small molecule / hetero: FF type (e.g. "oh", "c3") takes priority
            atom.force_field_type
                .clone()
                .or_else(|| atom.type_in_res.as_ref().map(|t| t.to_string()))
                .or_else(|| atom.type_in_res_general.clone())
                .unwrap_or_else(|| format!("{}{}", atom.element.to_letter(), nr))
        } else {
            // Protein / standard residue: residue atom name ("CA", "CB") takes priority
            atom.type_in_res
                .as_ref()
                .map(|t| t.to_string())
                .or_else(|| atom.force_field_type.clone())
                .or_else(|| atom.type_in_res_general.clone())
                .unwrap_or_else(|| format!("{}{}", atom.element.to_letter(), nr))
        };
        s.push_str(&format!(
            "  {:>4}  {:<6}  {:>5}  {:<8}  {:<6}  {:>4}  {:>8.4}  {:>8.3}\n",
            nr, ff_type, 1, mol.name, atom_name, nr, charge, mass,
        ));
    }
    s.push('\n');

    // [ bonds ]
    if !mol.bonds.is_empty() {
        s.push_str("[ bonds ]\n; ai   aj   funct  r0 (nm)    kb (kJ/mol/nm²)\n");
        for bond in mol.bonds {
            let (Some(&ai), Some(&aj)) = (
                sn_to_idx.get(&bond.atom_0_sn),
                sn_to_idx.get(&bond.atom_1_sn),
            ) else {
                continue;
            };
            let fft_i = mol
                .atoms
                .get(ai - 1)
                .and_then(|a| a.force_field_type.as_deref())
                .unwrap_or("X");
            let fft_j = mol
                .atoms
                .get(aj - 1)
                .and_then(|a| a.force_field_type.as_deref())
                .unwrap_or("X");
            let (r0_nm, kb_kj) = lookup_bond(fft_i, fft_j, ff, ff_global)
                .map(|(r0, kb)| (r0 * ANG_TO_NM, kb * BOND_K_FACTOR))
                .unwrap_or((0.1522, 224262.4)); // C-C fallback
            s.push_str(&format!(
                "  {:>4}  {:>4}  1  {:>12.6}  {:>12.1}\n",
                ai, aj, r0_nm, kb_kj,
            ));
        }
        s.push('\n');
    }

    // [ angles ]
    let angle_triples = enumerate_angles(&sn_to_idx, mol.bonds);
    if !angle_triples.is_empty() {
        s.push_str("[ angles ]\n; ai   aj   ak   funct  theta0 (deg)  ktheta (kJ/mol/rad²)\n");
        for (ai, aj, ak) in &angle_triples {
            let fft_i = mol
                .atoms
                .get(ai - 1)
                .and_then(|a| a.force_field_type.as_deref())
                .unwrap_or("X");
            let fft_j = mol
                .atoms
                .get(aj - 1)
                .and_then(|a| a.force_field_type.as_deref())
                .unwrap_or("X");
            let fft_k = mol
                .atoms
                .get(ak - 1)
                .and_then(|a| a.force_field_type.as_deref())
                .unwrap_or("X");
            let (theta_deg, k_kj) = lookup_angle(fft_i, fft_j, fft_k, ff, ff_global)
                .map(|(theta_rad, k)| (theta_rad.to_degrees(), k * KCAL_TO_KJ))
                .unwrap_or((109.5, 418.4));
            s.push_str(&format!(
                "  {:>4}  {:>4}  {:>4}  1  {:>10.3}  {:>12.3}\n",
                ai, aj, ak, theta_deg, k_kj,
            ));
        }
        s.push('\n');
    }

    // [ dihedrals ] — proper (funct 9) and improper (funct 4)
    let dihedral_quads = enumerate_dihedrals(&sn_to_idx, mol.bonds);
    if !dihedral_quads.is_empty() {
        s.push_str("[ dihedrals ]\n; ai   aj   ak   al   funct  phi0 (deg)  kphi (kJ/mol)  mult\n");
        for (ai, aj, ak, al) in &dihedral_quads {
            let fft_i = mol
                .atoms
                .get(ai - 1)
                .and_then(|a| a.force_field_type.as_deref())
                .unwrap_or("X");
            let fft_j = mol
                .atoms
                .get(aj - 1)
                .and_then(|a| a.force_field_type.as_deref())
                .unwrap_or("X");
            let fft_k = mol
                .atoms
                .get(ak - 1)
                .and_then(|a| a.force_field_type.as_deref())
                .unwrap_or("X");
            let fft_l = mol
                .atoms
                .get(al - 1)
                .and_then(|a| a.force_field_type.as_deref())
                .unwrap_or("X");
            if let Some(terms) = lookup_dihedral(fft_i, fft_j, fft_k, fft_l, ff, ff_global) {
                for t in terms {
                    s.push_str(&format!(
                        "  {:>4}  {:>4}  {:>4}  {:>4}  9  {:>10.3}  {:>12.4}  {:>4}\n",
                        ai,
                        aj,
                        ak,
                        al,
                        t.phase.to_degrees(),
                        t.barrier_height * KCAL_TO_KJ,
                        t.periodicity,
                    ));
                }
            }
        }
        s.push('\n');
    }

    // [ dihedrals ] — improper (funct 4)
    let improper_quads = enumerate_impropers(&sn_to_idx, mol.bonds);
    if !improper_quads.is_empty() {
        let mut improper_lines = String::new();
        for (ai, aj, ak, al) in &improper_quads {
            let fft_i = mol
                .atoms
                .get(ai - 1)
                .and_then(|a| a.force_field_type.as_deref())
                .unwrap_or("X");
            let fft_j = mol
                .atoms
                .get(aj - 1)
                .and_then(|a| a.force_field_type.as_deref())
                .unwrap_or("X");
            let fft_k = mol
                .atoms
                .get(ak - 1)
                .and_then(|a| a.force_field_type.as_deref())
                .unwrap_or("X");
            let fft_l = mol
                .atoms
                .get(al - 1)
                .and_then(|a| a.force_field_type.as_deref())
                .unwrap_or("X");
            if let Some(terms) = lookup_improper(fft_i, fft_j, fft_k, fft_l, ff, ff_global) {
                for t in terms {
                    improper_lines.push_str(&format!(
                        "  {:>4}  {:>4}  {:>4}  {:>4}  4  {:>10.3}  {:>12.4}  {:>4}\n",
                        ai,
                        aj,
                        ak,
                        al,
                        t.phase.to_degrees(),
                        t.barrier_height * KCAL_TO_KJ,
                        t.periodicity,
                    ));
                }
            }
        }
        if !improper_lines.is_empty() {
            s.push_str("[ dihedrals ] ; improper\n");
            s.push_str("; ai   aj   ak   al   funct  phi0 (deg)  kphi (kJ/mol)  mult\n");
            s.push_str(&improper_lines);
            s.push('\n');
        }
    }
}