bio_files 0.5.2

//! 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, io};

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

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>,
) -> io::Result<String> {
    let mut s = String::from("; GROMACS topology generated by Bio Files\n\n");

    // --- Pre-build indexed params for every molecule -------------------------
    // Do this first so [atomtypes] can read resolved mass/LJ without duplicating
    // alias logic — ForceFieldParamsIndexed::new is the single source of truth.
    let mol_params: Vec<ForceFieldParamsIndexed> = molecules
        .iter()
        .map(|mol| build_indexed(mol, ff_global))
        .collect::<io::Result<Vec<_>>>()?;

    // --- [ defaults ] -------------------------------------------------------
    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 ] -------------------------------------------------------
    // Read mass/LJ from the already-resolved indexed params — no alias logic here.
    // (ff_type → first atom we encounter that carries it)
    let mut atomtypes: HashMap<String, (f32, f32, f32)> = HashMap::new(); // ff_type → (mass, sigma, eps)
    for (mol, pi) in molecules.iter().zip(&mol_params) {
        for (i, atom) in mol.atoms.iter().enumerate() {
            if let Some(ff_type) = &atom.force_field_type {
                if !atomtypes.contains_key(ff_type) {
                    let mass = pi.mass.get(&i).map(|m| m.mass).unwrap_or(12.011);
                    let (sigma, eps) = pi
                        .lennard_jones
                        .get(&i)
                        .map(|lj| (lj.sigma, lj.eps))
                        .unwrap_or((0., 0.));
                    atomtypes.insert(ff_type.clone(), (mass, sigma, eps));
                }
            }
        }
    }

    s.push_str("[ atomtypes ]\n");
    s.push_str("; name  at.num   mass      charge  ptype  sigma (nm)      epsilon (kJ/mol)\n");

    let mut sorted_types: Vec<_> = atomtypes.iter().collect();
    sorted_types.sort_by_key(|(t, _)| t.as_str());
    for (ff_type, &(mass, sigma, eps)) in sorted_types {
        let at_num = atomic_number_from_mass(mass);
        s.push_str(&format!(
            "  {:<6}  {:>3}    {:>8.4}   0.000  A  {:>14.8e}  {:>14.8e}\n",
            ff_type,
            at_num,
            mass,
            sigma * ANG_TO_NM,
            eps * KCAL_TO_KJ,
        ));
    }

    // Water model atom types
    if let Some(WaterModel::Opc(_)) = water_model {
        s.push_str(&opc_atomtypes());
    }
    // Always include ion atom types — gmx genion may add NA/CL regardless of water model.
    s.push_str(&ion_atomtypes());
    s.push('\n');

    // --- Per-molecule blocks -------------------------------------------------
    for (mol, pi) in molecules.iter().zip(&mol_params) {
        write_molecule_block(&mut s, mol, pi)?;
    }

    // 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());
    }
    // Ion moleculetypes — gmx genion appends NA/CL to [molecules]; they must be defined here.
    s.push_str(&ion_moleculetypes());

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

    // --- [ molecules ] -------------------------------------------------------
    s.push_str("[ molecules ]\n; Compound     #mols\n");
    for mol in molecules {
        s.push_str(&format!("{:<14}  {}\n", mol.name, mol.count));
    }

    Ok(s)
}

/// Build the merged FF params and indexed params for a single molecule.
fn build_indexed(
    mol: &MoleculeTopology<'_>,
    ff_global: Option<&ForceFieldParams>,
) -> io::Result<ForceFieldParamsIndexed> {
    let merged;
    let ff: &ForceFieldParams = match (mol.ff_mol, ff_global) {
        (Some(mol_ff), Some(global)) => {
            merged = mol_ff.merge_with(global);
            &merged
        }
        (Some(mol_ff), None) => mol_ff,
        (None, Some(global)) => global,
        (None, None) => return Err(io::Error::other("Missing FF params for molecule")),
    };

    let sn_to_i: HashMap<u32, usize> = mol
        .atoms
        .iter()
        .enumerate()
        .map(|(i, a)| (a.serial_number, i))
        .collect();

    let mut adj_list: Vec<Vec<usize>> = vec![Vec::new(); mol.atoms.len()];
    for bond in mol.bonds {
        let (Some(&i0), Some(&i1)) = (sn_to_i.get(&bond.atom_0_sn), sn_to_i.get(&bond.atom_1_sn))
        else {
            continue;
        };
        adj_list[i0].push(i1);
        adj_list[i1].push(i0);
    }

    ForceFieldParamsIndexed::new(ff, mol.atoms, &adj_list, false)
}

// ---------------------------------------------------------------------------
// 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",
    ))
}

/// Atom-type lines for monovalent ions (Na⁺, Cl⁻) to append to `[atomtypes]`.
///
/// Parameters: Joung & Cheatham (2008) JPCB, widely used with Amber force fields in GROMACS.
/// sigma and epsilon in GROMACS units (nm, kJ/mol).
fn ion_atomtypes() -> &'static str {
    // Na+: R*=1.2195 Å → sigma=0.24393 nm; ε=0.01150 kJ/mol
    // Cl-: R*=2.2395 Å → sigma=0.44789 nm; ε=0.14890 kJ/mol
    concat!(
        "  Na+      11   22.98977   0.000  A   2.43928e-01  1.15897e-02\n",
        "  Cl-      17   35.45300   0.000  A   4.47766e-01  1.48913e-01\n",
    )
}

/// Moleculetype blocks for NA (sodium) and CL (chloride).
///
/// gmx genion uses these exact moleculetype names when adding counter-ions.
/// They must be present in the topology even if not initially in the system.
fn ion_moleculetypes() -> &'static str {
    concat!(
        "[ moleculetype ]\n",
        "; Name  nrexcl\n",
        "NA         1\n",
        "\n",
        "[ atoms ]\n",
        "; nr  type  resnr  residue  atom  cgnr  charge    mass\n",
        "  1   Na+   1      NA       NA    1     1.00000   22.98977\n",
        "\n",
        "[ moleculetype ]\n",
        "; Name  nrexcl\n",
        "CL         1\n",
        "\n",
        "[ atoms ]\n",
        "; nr  type  resnr  residue  atom  cgnr  charge    mass\n",
        "  1   Cl-   1      CL       CL    1    -1.00000   35.45300\n",
        "\n",
    )
}

// ---------------------------------------------------------------------------
// 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<'_>,
    pi: &ForceFieldParamsIndexed,
) -> io::Result<()> {
    // [ 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; // 1-based for GROMACS

        let Some(ff_type) = &atom.force_field_type else {
            return Err(io::Error::other("Missing FF type on atom; aborting"));
        };

        let charge = atom.partial_charge.unwrap_or(0.0);
        // Use indexed mass which already resolved aliases / element fallbacks.
        let mass = pi.mass.get(&i).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 {
            // todo: QC this logic of ff_type vs TIR general... Sus. They are NOT the same!
            // 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 ]
    // Iterate over the pre-built indexed params (0-based), emit 1-based for GROMACS.
    if !pi.bond_stretching.is_empty() {
        s.push_str("[ bonds ]\n; ai   aj   funct  r0 (nm)    kb (kJ/mol/nm²)\n");

        let mut bonds: Vec<_> = pi.bond_stretching.iter().collect();
        bonds.sort_by_key(|&(&(i0, i1), _)| (i0, i1));

        for (&(i0, i1), p) in bonds {
            s.push_str(&format!(
                "  {:>4}  {:>4}  1  {:>12.6e}  {:>12.1}\n",
                i0 + 1,
                i1 + 1,
                p.r_0 * ANG_TO_NM,
                p.k_b * BOND_K_FACTOR,
            ));
        }
        s.push('\n');
    }

    // [ angles ]
    if !pi.angle.is_empty() {
        s.push_str("[ angles ]\n; ai   aj   ak   funct  theta0 (deg)  ktheta (kJ/mol/rad²)\n");
        let mut angles: Vec<_> = pi.angle.iter().collect();
        angles.sort_by_key(|&(&(n0, ctr, n1), _)| (n0, ctr, n1));

        for (&(n0, ctr, n1), p) in angles {
            s.push_str(&format!(
                "  {:>4}  {:>4}  {:>4}  1  {:>10.3}  {:>12.3}\n",
                n0 + 1,
                ctr + 1,
                n1 + 1,
                p.theta_0.to_degrees(),
                p.k * KCAL_TO_KJ,
            ));
        }
        s.push('\n');
    }

    // [ dihedrals ] — proper (funct 9)
    if !pi.dihedral.is_empty() {
        s.push_str("[ dihedrals ]\n; ai   aj   ak   al   funct  phi0 (deg)  kphi (kJ/mol)  mult\n");
        let mut dihedrals: Vec<_> = pi.dihedral.iter().collect();
        dihedrals.sort_by_key(|&(&(i0, i1, i2, i3), _)| (i0, i1, i2, i3));
        for (&(i0, i1, i2, i3), terms) in dihedrals {
            for p in terms {
                s.push_str(&format!(
                    "  {:>4}  {:>4}  {:>4}  {:>4}  9  {:>10.3}  {:>12.4}  {:>4}\n",
                    i0 + 1,
                    i1 + 1,
                    i2 + 1,
                    i3 + 1,
                    p.phase.to_degrees(),
                    p.barrier_height * KCAL_TO_KJ,
                    p.periodicity,
                ));
            }
        }
        s.push('\n');
    }

    // [ dihedrals ] — improper (funct 4)
    if !pi.improper.is_empty() {
        s.push_str("[ dihedrals ] ; improper\n");
        s.push_str("; ai   aj   ak   al   funct  phi0 (deg)  kphi (kJ/mol)  mult\n");
        let mut impropers: Vec<_> = pi.improper.iter().collect();
        impropers.sort_by_key(|&(&(i0, i1, i2, i3), _)| (i0, i1, i2, i3));
        for (&(i0, i1, i2, i3), terms) in impropers {
            for p in terms {
                s.push_str(&format!(
                    "  {:>4}  {:>4}  {:>4}  {:>4}  4  {:>10.3}  {:>12.4}  {:>4}\n",
                    i0 + 1,
                    i1 + 1,
                    i2 + 1,
                    i3 + 1,
                    p.phase.to_degrees(),
                    p.barrier_height * KCAL_TO_KJ,
                    p.periodicity,
                ));
            }
        }
        s.push('\n');
    }

    Ok(())
}