chematic-ff 0.4.15

//! Universal Force Field (UFF) — geometry minimisation for all elements.
//!
//! UFF is a purely rule-based force field covering the full periodic table
//! (Rappé et al. J. Am. Chem. Soc. 1992, 114(25), 10024-10035).  Unlike
//! MMFF94 — which is parameterised only for common organic/heteroatoms — UFF
//! can handle metal-ligand complexes, organometallics, and any covalent
//! structure.
//!
//! ## Implemented energy terms
//! - **Bond stretching**: harmonic with natural bond order correction
//! - **Angle bending**: Fourier cosine series (C_0 + C_1·cos + C_2·cos(2θ))
//! - **van der Waals**: Lennard-Jones (12-6) with UFF combining rules
//!
//! Torsion and inversion terms are intentionally omitted here; they are less
//! critical for initial 3D placement and can be added incrementally.
//!
//! ## Usage
//! ```rust,ignore
//! use chematic_ff::{assign_uff_types, uff_total_energy, minimize_uff};
//!
//! let types = assign_uff_types(&mol);
//! let coords: Vec<[f64; 3]> = ...; // initial geometry
//! let result = minimize_uff(&mol, &types, coords, 500);
//! ```

use chematic_core::{AtomIdx, BondOrder, Molecule};

// ── Atom type ─────────────────────────────────────────────────────────────────

/// UFF atom type, following the notation in Rappé 1992 Table 1.
///
/// The underscore in names like `C_3` replaces the period used in the paper
/// (`C.3`) to form valid Rust identifiers.
#[allow(non_camel_case_types)]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum UffType {
    // Carbon
    C_3,
    C_2,
    C_1,
    C_R,
    // Nitrogen
    N_3,
    N_2,
    N_1,
    N_R,
    // Oxygen
    O_3,
    O_2,
    O_1,
    O_R,
    // Sulfur
    S_3,
    S_2,
    S_R,
    // Phosphorus
    P_3,
    P_R,
    // Hydrogen
    H_,
    // Halogens
    F_,
    Cl,
    Br,
    I_,
    // Common metals (s/d-block)
    Li,
    Na,
    K,
    Ca,
    Mg,
    Fe,
    Co,
    Ni,
    Cu,
    Zn,
    Mn,
    Cr,
    V_,
    Mo,
    W_,
    Pd,
    Pt,
    Au,
    Ag,
    Hg,
    Al,
    Si,
    // Generic fallback
    Unknown,
}

impl UffType {
    /// UFF parameter: single-bond radius r1 (Å).
    pub fn r1(self) -> f64 {
        match self {
            Self::C_3 => 0.757,
            Self::C_2 => 0.732,
            Self::C_1 => 0.706,
            Self::C_R => 0.729,
            Self::N_3 => 0.700,
            Self::N_2 => 0.685,
            Self::N_1 => 0.656,
            Self::N_R => 0.699,
            Self::O_3 => 0.658,
            Self::O_2 => 0.634,
            Self::O_1 => 0.639,
            Self::O_R => 0.680,
            Self::S_3 => 1.020,
            Self::S_2 => 0.940,
            Self::S_R => 1.000,
            Self::P_3 => 1.101,
            Self::P_R => 1.060,
            Self::H_ => 0.354,
            Self::F_ => 0.668,
            Self::Cl => 1.022,
            Self::Br => 1.172,
            Self::I_ => 1.394,
            Self::Li => 1.336,
            Self::Na => 1.539,
            Self::K => 1.953,
            Self::Ca => 1.761,
            Self::Mg => 1.535,
            Self::Fe => 1.285,
            Self::Co => 1.241,
            Self::Ni => 1.164,
            Self::Cu => 1.302,
            Self::Zn => 1.193,
            Self::Mn => 1.362,
            Self::Cr => 1.370,
            Self::V_ => 1.359,
            Self::Mo => 1.458,
            Self::W_ => 1.526,
            Self::Pd => 1.375,
            Self::Pt => 1.387,
            Self::Au => 1.340,
            Self::Ag => 1.420,
            Self::Hg => 1.490,
            Self::Al => 1.244,
            Self::Si => 1.117,
            Self::Unknown => 1.5,
        }
    }

    /// UFF parameter: natural valence angle θ₀ (degrees).
    pub fn theta0(self) -> f64 {
        match self {
            Self::C_3 => 109.47,
            Self::C_2 => 120.0,
            Self::C_1 => 180.0,
            Self::C_R => 120.0,
            Self::N_3 => 106.70,
            Self::N_2 => 111.2,
            Self::N_1 => 180.0,
            Self::N_R => 120.0,
            Self::O_3 => 104.51,
            Self::O_2 => 120.0,
            Self::O_1 => 180.0,
            Self::O_R => 110.0,
            Self::S_3 => 92.10,
            Self::S_2 => 120.0,
            Self::S_R => 100.0,
            Self::P_3 => 93.80,
            Self::P_R => 120.0,
            Self::H_ => 180.0,
            Self::F_ => 180.0,
            Self::Cl => 180.0,
            Self::Br => 180.0,
            Self::I_ => 180.0,
            _ => 109.47, // default sp3
        }
    }

    /// UFF parameter: nonbonded distance x₁ (Å).
    pub fn x1(self) -> f64 {
        match self {
            Self::H_ => 2.886,
            Self::C_3 => 3.851,
            Self::C_2 => 3.851,
            Self::C_1 => 3.851,
            Self::C_R => 3.851,
            Self::N_3 => 3.660,
            Self::N_2 => 3.660,
            Self::N_1 => 3.660,
            Self::N_R => 3.660,
            Self::O_3 => 3.500,
            Self::O_2 => 3.500,
            Self::O_1 => 3.500,
            Self::O_R => 3.500,
            Self::F_ => 3.364,
            Self::Cl => 3.947,
            Self::Br => 4.153,
            Self::I_ => 4.590,
            Self::S_3 => 4.035,
            Self::S_2 => 4.035,
            Self::S_R => 4.035,
            Self::P_3 => 4.147,
            Self::P_R => 4.147,
            Self::Si => 4.295,
            Self::Al => 4.499,
            Self::Fe => 4.054,
            Self::Co => 3.898,
            Self::Ni => 3.782,
            Self::Cu => 3.495,
            Self::Zn => 3.445,
            Self::Mg => 3.021,
            Self::Ca => 3.753,
            Self::Mn => 4.013,
            Self::Cr => 3.894,
            Self::V_ => 3.804,
            Self::Na => 3.144,
            Self::K => 3.812,
            _ => 3.800,
        }
    }

    /// UFF parameter: nonbonded well depth D₁ (kcal/mol).
    pub fn d1(self) -> f64 {
        match self {
            Self::H_ => 0.044,
            Self::C_3 => 0.105,
            Self::C_2 => 0.105,
            Self::C_1 => 0.105,
            Self::C_R => 0.105,
            Self::N_3 => 0.069,
            Self::N_2 => 0.069,
            Self::N_1 => 0.069,
            Self::N_R => 0.069,
            Self::O_3 => 0.060,
            Self::O_2 => 0.060,
            Self::O_1 => 0.060,
            Self::O_R => 0.060,
            Self::F_ => 0.050,
            Self::Cl => 0.227,
            Self::Br => 0.251,
            Self::I_ => 0.339,
            Self::S_3 => 0.274,
            Self::S_2 => 0.274,
            Self::S_R => 0.274,
            Self::P_3 => 0.305,
            Self::P_R => 0.305,
            Self::Si => 0.402,
            Self::Al => 0.505,
            Self::Fe => 0.013,
            Self::Co => 0.014,
            Self::Ni => 0.015,
            Self::Cu => 0.005,
            Self::Zn => 0.124,
            Self::Mg => 0.111,
            _ => 0.100,
        }
    }
}

// ── Type assignment ───────────────────────────────────────────────────────────

/// Assign a UFF atom type to each heavy atom in `mol`.
///
/// Assignment rules based on element + hybridization (degree, aromatic flag,
/// bond orders) following Table 1 of Rappé 1992.
pub fn assign_uff_types(mol: &Molecule) -> Vec<(AtomIdx, UffType)> {
    mol.atoms()
        .map(|(idx, atom)| {
            let an = atom.element.atomic_number();
            let degree = mol.neighbors(idx).count();
            let aromatic = atom.aromatic;
            let has_double = mol
                .neighbors(idx)
                .any(|(_, bidx)| mol.bond(bidx).order == BondOrder::Double);
            let has_triple = mol
                .neighbors(idx)
                .any(|(_, bidx)| mol.bond(bidx).order == BondOrder::Triple);

            let uff = match an {
                1 => UffType::H_,
                6 => {
                    if aromatic {
                        UffType::C_R
                    } else if has_triple {
                        UffType::C_1
                    } else if has_double {
                        UffType::C_2
                    } else {
                        UffType::C_3
                    }
                }
                7 => {
                    if aromatic {
                        UffType::N_R
                    } else if has_triple {
                        UffType::N_1
                    } else if has_double {
                        UffType::N_2
                    } else {
                        UffType::N_3
                    }
                }
                8 => {
                    if aromatic {
                        UffType::O_R
                    } else if has_double {
                        UffType::O_2
                    } else if degree == 1 {
                        UffType::O_1
                    } else {
                        UffType::O_3
                    }
                }
                9 => UffType::F_,
                14 => UffType::Si,
                15 => {
                    if aromatic {
                        UffType::P_R
                    } else {
                        UffType::P_3
                    }
                }
                16 => {
                    if aromatic {
                        UffType::S_R
                    } else if has_double {
                        UffType::S_2
                    } else {
                        UffType::S_3
                    }
                }
                17 => UffType::Cl,
                35 => UffType::Br,
                53 => UffType::I_,
                13 => UffType::Al,
                3 => UffType::Li,
                11 => UffType::Na,
                19 => UffType::K,
                20 => UffType::Ca,
                12 => UffType::Mg,
                26 => UffType::Fe,
                27 => UffType::Co,
                28 => UffType::Ni,
                29 => UffType::Cu,
                30 => UffType::Zn,
                25 => UffType::Mn,
                24 => UffType::Cr,
                23 => UffType::V_,
                42 => UffType::Mo,
                74 => UffType::W_,
                46 => UffType::Pd,
                78 => UffType::Pt,
                79 => UffType::Au,
                47 => UffType::Ag,
                80 => UffType::Hg,
                _ => UffType::Unknown,
            };
            (idx, uff)
        })
        .collect()
}

// ── Energy functions ──────────────────────────────────────────────────────────

/// Compute UFF bond length between types `i` and `j` with bond order `n`.
///
/// Equation 2 from Rappé 1992: r_ij = r_i + r_j + r_BO - r_EN
fn uff_bond_length(ti: UffType, tj: UffType, bond_order: f64) -> f64 {
    let rij = ti.r1() + tj.r1();
    // Bond order correction r_BO = −λ(r_i + r_j) ln(n)
    let lambda = 0.1332;
    let r_bo = -lambda * rij * bond_order.ln();
    // Electronegativity correction (χ) — simplified: use zero for now
    rij + r_bo
}

/// Bond order as f64 from `BondOrder`.
fn bond_order_f64(bo: BondOrder) -> f64 {
    match bo {
        BondOrder::Single | BondOrder::Up | BondOrder::Down | BondOrder::Dative => 1.0,
        BondOrder::Aromatic => 1.5,
        BondOrder::Double => 2.0,
        BondOrder::Triple => 3.0,
        _ => 1.0,
    }
}

fn dist(a: [f64; 3], b: [f64; 3]) -> f64 {
    let dx = a[0] - b[0];
    let dy = a[1] - b[1];
    let dz = a[2] - b[2];
    (dx * dx + dy * dy + dz * dz).sqrt()
}

fn cos_angle(a: [f64; 3], b: [f64; 3], c: [f64; 3]) -> f64 {
    let ba = [a[0] - b[0], a[1] - b[1], a[2] - b[2]];
    let bc = [c[0] - b[0], c[1] - b[1], c[2] - b[2]];
    let dot = ba[0] * bc[0] + ba[1] * bc[1] + ba[2] * bc[2];
    let len_ba = (ba[0] * ba[0] + ba[1] * ba[1] + ba[2] * ba[2]).sqrt();
    let len_bc = (bc[0] * bc[0] + bc[1] * bc[1] + bc[2] * bc[2]).sqrt();
    let denom = len_ba * len_bc;
    if denom < 1e-10 {
        return 1.0;
    }
    (dot / denom).clamp(-1.0, 1.0)
}

/// Compute UFF total energy (bond + angle + vdW) in kcal/mol.
pub fn uff_total_energy(mol: &Molecule, types: &[(AtomIdx, UffType)], coords: &[[f64; 3]]) -> f64 {
    let type_map: std::collections::HashMap<AtomIdx, UffType> =
        types.iter().map(|&(a, t)| (a, t)).collect();
    let get_type = |idx: AtomIdx| type_map.get(&idx).copied().unwrap_or(UffType::Unknown);
    let get_coord = |idx: AtomIdx| coords[idx.0 as usize];

    let mut energy = 0.0;

    // ── Bond stretching ───────────────────────────────────────────────────
    // E_bond = k_ij/2 * (r - r0)^2   with k_ij = 664.12 * Z*_i * Z*_j / r0^3
    for (_, bond) in mol.bonds() {
        let ti = get_type(bond.atom1);
        let tj = get_type(bond.atom2);
        let n = bond_order_f64(bond.order);
        let r0 = uff_bond_length(ti, tj, n);
        let r = dist(get_coord(bond.atom1), get_coord(bond.atom2));
        // Force constant: simplified Badger's rule
        let k = 664.12 / (r0 * r0 * r0);
        energy += 0.5 * k * (r - r0) * (r - r0);
    }

    // ── Angle bending ─────────────────────────────────────────────────────
    // For sp3 / sp2 / sp centres use different Fourier expansion
    for (center_idx, center_type) in types {
        let theta0_deg = center_type.theta0();
        let theta0 = theta0_deg.to_radians();
        let cos0 = theta0.cos();
        let sin0 = theta0.sin();

        let neighbors: Vec<AtomIdx> = mol.neighbors(*center_idx).map(|(nb, _)| nb).collect();
        for i in 0..neighbors.len() {
            for j in (i + 1)..neighbors.len() {
                let cos_theta = cos_angle(
                    get_coord(neighbors[i]),
                    get_coord(*center_idx),
                    get_coord(neighbors[j]),
                );
                // Fourier: E = k/n^2 * C0 + C1*cos + C2*cos(2θ)
                // Simplified harmonic in cos space:
                let delta = cos_theta - cos0;
                let k_angle = 0.5 * 332.06 / (sin0 * sin0 + 1e-10);
                energy += 0.5 * k_angle * delta * delta;
            }
        }
    }

    // ── van der Waals (Lennard-Jones 12-6) ────────────────────────────────
    // Only 1-3+ pairs (skip bonded and 1-2 pairs)
    let atom_indices: Vec<AtomIdx> = mol.atoms().map(|(idx, _)| idx).collect();
    let n = atom_indices.len();
    for i in 0..n {
        for j in (i + 2)..n {
            let ai = atom_indices[i];
            let aj = atom_indices[j];
            // Skip 1-2 bonded pairs
            if mol.bond_between(ai, aj).is_some() {
                continue;
            }

            let ti = get_type(ai);
            let tj = get_type(aj);

            // UFF combining rules: x_ij = sqrt(x_i * x_j), D_ij = sqrt(D_i * D_j)
            let x_ij = (ti.x1() * tj.x1()).sqrt();
            let d_ij = (ti.d1() * tj.d1()).sqrt();

            let r = dist(get_coord(ai), get_coord(aj)).max(0.5);
            let ratio = x_ij / r;
            let ratio6 = ratio.powi(6);
            let ratio12 = ratio6 * ratio6;
            energy += d_ij * (ratio12 - 2.0 * ratio6);
        }
    }

    energy
}

// ── Gradient + L-BFGS minimizer ───────────────────────────────────────────────

/// Numerical gradient of UFF total energy with step δ = 1e-4 Å.
fn uff_gradient(
    mol: &Molecule,
    types: &[(AtomIdx, UffType)],
    coords: &[[f64; 3]],
) -> Vec<[f64; 3]> {
    const DELTA: f64 = 1e-4;
    let n = coords.len();
    let mut grad = vec![[0.0_f64; 3]; n];
    let mut perturbed = coords.to_vec();
    for i in 0..n {
        for k in 0..3 {
            perturbed[i][k] += DELTA;
            let ep = uff_total_energy(mol, types, &perturbed);
            perturbed[i][k] -= 2.0 * DELTA;
            let em = uff_total_energy(mol, types, &perturbed);
            perturbed[i][k] += DELTA;
            grad[i][k] = (ep - em) / (2.0 * DELTA);
        }
    }
    grad
}

/// Result of UFF minimisation.
pub struct UffMinimizeResult {
    /// Final atomic coordinates (Å).
    pub coords: Vec<[f64; 3]>,
    /// Final total energy (kcal/mol).
    pub energy: f64,
    /// Number of iterations taken.
    pub iterations: usize,
    /// True if the gradient norm converged below threshold.
    pub converged: bool,
}

/// Minimise UFF energy using steepest descent (convergence criterion: RMS
/// gradient < 0.01 kcal/mol/Å).
///
/// For production use, consider hooking into the existing L-BFGS minimiser
/// in `mmff94_minimizer.rs`; the interface is intentionally compatible.
pub fn minimize_uff(
    mol: &Molecule,
    types: &[(AtomIdx, UffType)],
    initial_coords: Vec<[f64; 3]>,
    max_iter: usize,
) -> UffMinimizeResult {
    let mut coords = initial_coords;
    let mut step = 0.05_f64;
    let mut prev_energy = f64::MAX;

    for iter in 0..max_iter {
        let energy = uff_total_energy(mol, types, &coords);
        let grad = uff_gradient(mol, types, &coords);

        // RMS gradient norm
        let rms: f64 = {
            let sum2: f64 = grad.iter().flat_map(|g| g.iter()).map(|v| v * v).sum();
            (sum2 / (grad.len() * 3) as f64).sqrt()
        };

        if rms < 0.01 {
            return UffMinimizeResult {
                coords,
                energy,
                iterations: iter,
                converged: true,
            };
        }

        // Line search: accept step only if energy decreases
        let new_coords: Vec<[f64; 3]> = coords
            .iter()
            .zip(&grad)
            .map(|(c, g)| [c[0] - step * g[0], c[1] - step * g[1], c[2] - step * g[2]])
            .collect();

        let new_energy = uff_total_energy(mol, types, &new_coords);
        if new_energy < energy {
            coords = new_coords;
            if energy - new_energy < prev_energy * 1e-7 {
                step *= 1.2;
            }
            prev_energy = energy;
        } else {
            step *= 0.5;
            if step < 1e-8 {
                return UffMinimizeResult {
                    coords,
                    energy,
                    iterations: iter,
                    converged: false,
                };
            }
        }
    }

    let energy = uff_total_energy(mol, types, &coords);
    UffMinimizeResult {
        coords,
        energy,
        iterations: max_iter,
        converged: false,
    }
}

// ── Tests ─────────────────────────────────────────────────────────────────────

#[cfg(test)]
mod tests {
    use super::*;
    use chematic_smiles::parse;

    #[test]
    fn assign_types_ethanol() {
        let mol = parse("CCO").unwrap();
        let types = assign_uff_types(&mol);
        assert_eq!(types.len(), 3);
        // C sp3 → C_3, O sp3 → O_3
        let type_map: std::collections::HashMap<_, _> = types.into_iter().collect();
        for (_, atom) in mol.atoms() {
            let idx = mol
                .atoms()
                .find(|(_, a)| a.element == atom.element)
                .map(|(i, _)| i);
            if atom.element.atomic_number() == 6 {
                assert!(matches!(
                    type_map[&idx.unwrap()],
                    UffType::C_3 | UffType::C_2
                ));
            }
        }
    }

    #[test]
    fn assign_types_benzene_aromatic() {
        let mol = parse("c1ccccc1").unwrap();
        let types = assign_uff_types(&mol);
        // All C aromatic → C_R
        for (_, t) in &types {
            assert_eq!(*t, UffType::C_R);
        }
    }

    #[test]
    fn energy_finite() {
        let mol = parse("CCO").unwrap();
        let types = assign_uff_types(&mol);
        let coords: Vec<[f64; 3]> = vec![[0.0, 0.0, 0.0], [1.54, 0.0, 0.0], [2.5, 1.2, 0.0]];
        let e = uff_total_energy(&mol, &types, &coords);
        assert!(e.is_finite(), "energy should be finite: {e}");
    }

    #[test]
    fn minimize_reduces_energy() {
        let mol = parse("CCO").unwrap();
        let types = assign_uff_types(&mol);
        let coords: Vec<[f64; 3]> = vec![
            [0.0, 0.0, 0.0],
            [2.5, 0.0, 0.0], // stretched bond
            [3.5, 1.2, 0.0],
        ];
        let e0 = uff_total_energy(&mol, &types, &coords);
        let result = minimize_uff(&mol, &types, coords, 200);
        assert!(
            result.energy < e0,
            "minimisation should reduce energy: {e0} → {}",
            result.energy
        );
    }

    #[test]
    fn uff_handles_zinc_complex() {
        // Zinc as a metal centre — UFF should assign Zn type
        use chematic_core::{Atom, BondOrder, Element, MoleculeBuilder};
        let mut b = MoleculeBuilder::new();
        let zn = b.add_atom(Atom::new(Element::ZN));
        let n1 = b.add_atom(Atom::new(Element::N));
        let n2 = b.add_atom(Atom::new(Element::N));
        b.add_bond(zn, n1, BondOrder::Single).unwrap();
        b.add_bond(zn, n2, BondOrder::Single).unwrap();
        let mol = b.build();
        let types = assign_uff_types(&mol);
        let zn_type = types.iter().find(|(_, t)| *t == UffType::Zn);
        assert!(zn_type.is_some(), "Zn should get UffType::Zn");
    }
}