chemx-props 0.2.0

//! Rigid-rotor / harmonic-oscillator (RRHO) thermochemistry at 298.15 K.
//!
//! All energies are in hartree per molecule. Entropy is in hartree/(K·molecule).
//! Multiply energies by `HARTREE_TO_KCAL_MOL` or `HARTREE_TO_KJ_MOL` for
//! molar units; S by the same factor for molar entropy in kcal/(mol·K).

use std::f64::consts::PI;

use chemx_core::{Molecule, units::BOLTZMANN_HT};
use chemx_linalg::symmetric_eigh;

use crate::frequencies::FrequencyResult;

// SI physical constants (CODATA 2018 / 2019 SI redefinitions)
const KB_J: f64 = 1.380_649e-23; // J/K (exact)
const H_J: f64 = 6.626_070_15e-34; // J·s (exact)
const C_CM: f64 = 2.997_924_58e10; // cm/s (exact)
const AMU_KG: f64 = 1.660_539_066_60e-27; // kg/amu (CODATA 2018)
const BOHR_M: f64 = 5.291_772_109_03e-11; // m/bohr (CODATA 2018)
const EH_J: f64 = 4.359_744_722_207_1e-18; // J/Eh (CODATA 2018)
const STD_P_PA: f64 = 101_325.0; // Pa (1 atm)

/// RRHO thermochemistry result at temperature `T`.
#[derive(Debug, Clone)]
pub struct ThermoResult {
    /// Temperature, K.
    pub temperature: f64,
    /// Rotational symmetry number used.
    pub symmetry_number: u32,
    /// Zero-point vibrational energy (hartree).
    pub zpe: f64,
    /// Thermal correction to internal energy = ZPE + ΔU_vib + U_rot + U_trans.
    pub thermal_energy_corr: f64,
    /// Thermal correction to enthalpy = thermal_energy_corr + k_B T (PV = RT for
    /// ideal gas).
    pub enthalpy_corr: f64,
    /// Total enthalpy H(T) = E_elec + enthalpy_corr (hartree).
    pub enthalpy: f64,
    /// Total entropy (hartree/K).
    pub entropy: f64,
    /// Gibbs free energy G(T) = H(T) − T·S (hartree).
    pub gibbs: f64,
    /// Vibrational frequencies used (cm⁻¹); only the positive (real) ones enter
    /// the vibrational partition function.
    pub vib_frequencies_cm1: Vec<f64>,
    /// Principal moments of inertia in amu·bohr² [I_A ≤ I_B ≤ I_C].
    pub moments_of_inertia: Vec<f64>,
    /// Whether the molecule is linear (one moment < threshold).
    pub is_linear: bool,
}

/// Compute RRHO thermochemistry from a [`FrequencyResult`] and the electronic energy.
///
/// `electronic_energy`: total SCF/post-HF energy in hartree.
/// `temperature`: K (298.15 for standard conditions).
/// `symmetry_number`: rotational symmetry number σ (1 for C₁; 2 for H₂O, HCl;
///   12 for benzene, CH₄). A wrong σ shifts entropy and Gibbs by k_B·ln(σ) per
///   molecule — the usual source of thermo discrepancy with ORCA.
/// `multiplicity`: electronic-state multiplicity (contributes S_elec = k_B ln(mult)).
///
/// The vibrational partition function uses the 3N−6 (3N−5 if linear) genuine
/// vibrations: the trans/rot modes (the 6/5 lowest after Eckart projection) are
/// dropped by count, not by a `freq > 0` filter — a trans/rot mode that projects
/// to a tiny positive residue would otherwise diverge S_vib and U_vib.
pub fn rrho_thermochemistry(
    molecule: &Molecule,
    freq_result: &FrequencyResult,
    electronic_energy: f64,
    temperature: f64,
    symmetry_number: u32,
    multiplicity: u32,
) -> ThermoResult {
    let kt_j = KB_J * temperature; // k_B T in Joules
    let kt_eh = BOLTZMANN_HT * temperature; // k_B T in hartree

    // Principal moments of inertia
    let (moments, is_linear) = principal_moments(molecule);

    // Total molecular mass in kg
    let total_mass_kg: f64 = molecule
        .atoms
        .iter()
        .map(|a| a.element.mass() * AMU_KG)
        .sum();

    // ─── Translational ───────────────────────────────────────────────────────
    // q_trans = (2π m k_B T)^(3/2) k_B T / (h³ P°)  [per molecule]
    let q_trans = (2.0 * PI * total_mass_kg * kt_j).powf(1.5) * kt_j / (H_J.powi(3) * STD_P_PA);
    let s_trans_j = KB_J * (2.5 + q_trans.ln()); // Sackur–Tetrode
    let u_trans_j = 1.5 * kt_j;

    // ─── Rotational ──────────────────────────────────────────────────────────
    let (_q_rot, s_rot_j, u_rot_j) = if is_linear {
        // One moment of inertia (I_C = largest; I_A = I_B; I_A ≈ 0 for linear)
        let i_si = moments[2] * AMU_KG * BOHR_M * BOHR_M; // kg·m²
        let q = 8.0 * PI * PI * i_si * kt_j / (H_J * H_J * symmetry_number as f64);
        let s = KB_J * (1.0 + q.ln());
        let u = kt_j;
        (q, s, u)
    } else {
        // Three non-zero moments
        let i_si: Vec<f64> = moments
            .iter()
            .map(|m| m * AMU_KG * BOHR_M * BOHR_M)
            .collect();
        let prod_i = i_si[0] * i_si[1] * i_si[2];
        let q = PI.sqrt() / symmetry_number as f64
            * (8.0 * PI * PI * kt_j / (H_J * H_J)).powf(1.5)
            * prod_i.sqrt();
        let s = KB_J * (1.5 + q.ln());
        let u = 1.5 * kt_j;
        (q, s, u)
    };

    // ─── Vibrational ─────────────────────────────────────────────────────────
    // The vibrational partition function runs over the 3N−6 (3N−5 if linear)
    // genuine vibrations. `harmonic_frequencies` returns all 3N modes sorted
    // ascending with the trans/rot modes (≈0 after Eckart projection) first, so
    // drop exactly those by count. A bare `f > 0` filter is NOT safe: a trans/rot
    // mode that projects to a tiny *positive* residue (≈1e-6 cm⁻¹) would slip
    // through and blow up the entropy (S_vib ∝ −ln(1−e^{−hν/kT}) → ∞ as ν→0) and
    // the thermal energy (∝ 1/(e^{hν/kT}−1) → ∞).
    let n_trans_rot = if is_linear { 5 } else { 6 };
    let vib_cm1: Vec<f64> = freq_result
        .frequencies_cm1
        .iter()
        .skip(n_trans_rot)
        .copied()
        .filter(|&f| f > 0.0) // defensive: ignore any residual imaginary mode
        .collect();

    let mut zpe_j = 0.0f64;
    let mut u_vib_thermal_j = 0.0f64;
    let mut s_vib_j = 0.0f64;
    for &nu_cm1 in &vib_cm1 {
        let nu_hz = nu_cm1 * C_CM; // Hz
        let x = H_J * nu_hz / kt_j; // dimensionless h ν / k_B T
        zpe_j += 0.5 * H_J * nu_hz;
        u_vib_thermal_j += H_J * nu_hz / (x.exp() - 1.0);
        s_vib_j += KB_J * (x / (x.exp() - 1.0) - (1.0 - (-x).exp()).ln());
    }

    // ─── Electronic ──────────────────────────────────────────────────────────
    let s_elec_j = KB_J * (multiplicity as f64).ln();

    // ─── Assemble ────────────────────────────────────────────────────────────
    // Convert to hartree
    let to_eh = |j: f64| j / EH_J;
    let zpe = to_eh(zpe_j);
    let u_vib_thermal = to_eh(u_vib_thermal_j);
    let u_rot = to_eh(u_rot_j);
    let u_trans = to_eh(u_trans_j);
    // The s_*_j are already entropies (J/K, = k_B · dimensionless); converting
    // J/K → Eh/K is a single division by EH_J. (Do NOT also divide by T — these
    // are S, not the T·S "entropy term" some codes print.)
    let s_vib = to_eh(s_vib_j);
    let s_rot = to_eh(s_rot_j);
    let s_trans = to_eh(s_trans_j);
    let s_elec = to_eh(s_elec_j);

    let thermal_energy_corr = zpe + u_vib_thermal + u_rot + u_trans;
    let enthalpy_corr = thermal_energy_corr + kt_eh; // PV = k_B T for ideal gas
    let enthalpy = electronic_energy + enthalpy_corr;
    let entropy = s_vib + s_rot + s_trans + s_elec;
    let gibbs = enthalpy - temperature * entropy;

    ThermoResult {
        temperature,
        symmetry_number,
        zpe,
        thermal_energy_corr,
        enthalpy_corr,
        enthalpy,
        entropy,
        gibbs,
        vib_frequencies_cm1: vib_cm1,
        moments_of_inertia: moments,
        is_linear,
    }
}

/// Return principal moments of inertia [I_A ≤ I_B ≤ I_C] in amu·bohr², and
/// whether the molecule is linear (I_A < 1e-4 amu·bohr²).
fn principal_moments(molecule: &Molecule) -> (Vec<f64>, bool) {
    let masses: Vec<f64> = molecule.atoms.iter().map(|a| a.element.mass()).collect();
    let total_mass: f64 = masses.iter().sum();

    // Center of mass
    let mut com = [0.0f64; 3];
    for (i, atom) in molecule.atoms.iter().enumerate() {
        for (k, c) in com.iter_mut().enumerate() {
            *c += masses[i] * atom.position[k];
        }
    }
    for c in &mut com {
        *c /= total_mass;
    }

    // Inertia tensor (3×3)
    let mut imat = [[0.0f64; 3]; 3];
    for (i, atom) in molecule.atoms.iter().enumerate() {
        let r = [
            atom.position[0] - com[0],
            atom.position[1] - com[1],
            atom.position[2] - com[2],
        ];
        let r2 = r[0] * r[0] + r[1] * r[1] + r[2] * r[2];
        let m = masses[i];
        for a in 0..3 {
            for b in 0..3 {
                let delta = if a == b { 1.0 } else { 0.0 };
                imat[a][b] += m * (delta * r2 - r[a] * r[b]);
            }
        }
    }

    // Diagonalize the 3×3 inertia tensor via symmetric_eigh
    let flat: Vec<f64> = (0..3)
        .flat_map(|i| (0..3).map(move |j| imat[i][j]))
        .collect();
    let imat_fmat = chemx_linalg::mat_from_row_major(3, &flat);
    let eigh = symmetric_eigh(&imat_fmat);

    // Eigenvalues are moments in ascending order (amu·bohr²)
    let moments: Vec<f64> = eigh.values.iter().map(|&v| v.max(0.0)).collect();
    let is_linear = moments[0] < 1e-4;

    (moments, is_linear)
}