readcon_core/

units.rs

//! Unit expressions and conversion (metatomic-inspired, SI dimensional analysis).
//!
//! Supports named base units and simple products/quotients with `*`, `/`, `^`,
//! and parentheses (case-insensitive). `unit_conversion_factor(from, to)` returns
//! the multiplier `x_to = factor * x_from` when dimensions match.

use crate::error::ParseError;
use std::collections::HashMap;

/// Physical dimension exponents: L, T, M, Q (charge), Θ (temperature).
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct Dimension {
    pub exponents: [f64; 5],
}

impl Dimension {
    pub const ZERO: Self = Self {
        exponents: [0.0; 5],
    };
    pub const LENGTH: Self = Self {
        exponents: [1.0, 0.0, 0.0, 0.0, 0.0],
    };
    pub const TIME: Self = Self {
        exponents: [0.0, 1.0, 0.0, 0.0, 0.0],
    };
    pub const MASS: Self = Self {
        exponents: [0.0, 0.0, 1.0, 0.0, 0.0],
    };
    pub const CHARGE: Self = Self {
        exponents: [0.0, 0.0, 0.0, 1.0, 0.0],
    };

    fn mul(self, other: Self) -> Self {
        let mut e = [0.0; 5];
        for i in 0..5 {
            e[i] = self.exponents[i] + other.exponents[i];
        }
        Self { exponents: e }
    }

    fn div(self, other: Self) -> Self {
        let mut e = [0.0; 5];
        for i in 0..5 {
            e[i] = self.exponents[i] - other.exponents[i];
        }
        Self { exponents: e }
    }

    fn pow(self, p: f64) -> Self {
        let mut e = [0.0; 5];
        for i in 0..5 {
            e[i] = self.exponents[i] * p;
        }
        Self { exponents: e }
    }

    fn compatible(self, other: Self) -> bool {
        self.exponents
            .iter()
            .zip(other.exponents.iter())
            .all(|(a, b)| (a - b).abs() < 1e-9)
    }
}

/// Parsed unit: SI scale factor (value in SI = factor * value_in_unit) and dimensions.
#[derive(Clone, Debug)]
struct UnitValue {
    /// Multiply quantity in this unit by `si_factor` to get SI.
    si_factor: f64,
    dim: Dimension,
}

fn base_table() -> HashMap<&'static str, UnitValue> {
    let mut m = HashMap::new();
    let ins = |m: &mut HashMap<&'static str, UnitValue>, names: &[&'static str], u: UnitValue| {
        for n in names {
            m.insert(*n, u.clone());
        }
    };
    // Length → meters
    ins(
        &mut m,
        &["m", "meter", "metre"],
        UnitValue {
            si_factor: 1.0,
            dim: Dimension::LENGTH,
        },
    );
    ins(
        &mut m,
        &["angstrom", "a", "å"],
        UnitValue {
            si_factor: 1e-10,
            dim: Dimension::LENGTH,
        },
    );
    ins(
        &mut m,
        &["nm", "nanometer"],
        UnitValue {
            si_factor: 1e-9,
            dim: Dimension::LENGTH,
        },
    );
    ins(
        &mut m,
        &["bohr", "a0"],
        UnitValue {
            si_factor: 5.291_772_109_03e-11,
            dim: Dimension::LENGTH,
        },
    );
    // Time → seconds
    ins(
        &mut m,
        &["s", "second"],
        UnitValue {
            si_factor: 1.0,
            dim: Dimension::TIME,
        },
    );
    ins(
        &mut m,
        &["fs", "femtosecond"],
        UnitValue {
            si_factor: 1e-15,
            dim: Dimension::TIME,
        },
    );
    ins(
        &mut m,
        &["ps", "picosecond"],
        UnitValue {
            si_factor: 1e-12,
            dim: Dimension::TIME,
        },
    );
    ins(
        &mut m,
        &["ns", "nanosecond"],
        UnitValue {
            si_factor: 1e-9,
            dim: Dimension::TIME,
        },
    );
    // Mass → kg
    ins(
        &mut m,
        &["kg", "kilogram"],
        UnitValue {
            si_factor: 1.0,
            dim: Dimension::MASS,
        },
    );
    ins(
        &mut m,
        &["u", "amu", "dalton", "da"],
        UnitValue {
            si_factor: 1.660_539_066_60e-27,
            dim: Dimension::MASS,
        },
    );
    // Energy → joule (M L^2 T^-2)
    let energy_dim = Dimension::MASS
        .mul(Dimension::LENGTH.pow(2.0))
        .div(Dimension::TIME.pow(2.0));
    ins(
        &mut m,
        &["j", "joule"],
        UnitValue {
            si_factor: 1.0,
            dim: energy_dim,
        },
    );
    ins(
        &mut m,
        &["ev"],
        UnitValue {
            si_factor: 1.602_176_634e-19,
            dim: energy_dim,
        },
    );
    ins(
        &mut m,
        &["mev"],
        UnitValue {
            si_factor: 1.602_176_634e-22,
            dim: energy_dim,
        },
    );
    ins(
        &mut m,
        &["hartree", "ha"],
        UnitValue {
            si_factor: 4.359_744_722_207_1e-18,
            dim: energy_dim,
        },
    );
    // Charge
    ins(
        &mut m,
        &["e", "electron_charge"],
        UnitValue {
            si_factor: 1.602_176_634e-19,
            dim: Dimension::CHARGE,
        },
    );
    // 1 mol = N_A entities. `kcal / mol` is then 4184/N_A J per entity
    // (metatomic: eV -> kcal/mol is ~23.06).
    ins(
        &mut m,
        &["mol"],
        UnitValue {
            si_factor: 6.022_140_76e23,
            dim: Dimension::ZERO,
        },
    );
    ins(
        &mut m,
        &["kcal"],
        UnitValue {
            si_factor: 4184.0,
            dim: energy_dim,
        },
    );
    ins(
        &mut m,
        &["kj"],
        UnitValue {
            si_factor: 1000.0,
            dim: energy_dim,
        },
    );
    m
}

/// Parse a unit expression into SI factor and dimension.
pub fn parse_unit_expression(expr: &str) -> Result<(f64, Dimension), ParseError> {
    let s: String = expr.chars().filter(|c| !c.is_whitespace()).collect();
    if s.is_empty() {
        return Err(ParseError::ValidationError("empty unit expression".into()));
    }
    let table = base_table();
    parse_expr(&s.to_ascii_lowercase(), &table)
}

fn parse_expr(s: &str, table: &HashMap<&str, UnitValue>) -> Result<(f64, Dimension), ParseError> {
    // Split on * and / with left-associative scan; handle ^N on atoms.
    let mut factor = 1.0_f64;
    let mut dim = Dimension::ZERO;
    let mut i = 0;
    let bytes = s.as_bytes();
    let mut pending_div = false;
    while i < bytes.len() {
        if bytes[i] == b'*' {
            i += 1;
            continue;
        }
        if bytes[i] == b'/' {
            pending_div = true;
            i += 1;
            continue;
        }
        let (atom_f, atom_d, consumed) = parse_atom(&s[i..], table)?;
        i += consumed;
        if pending_div {
            factor /= atom_f;
            dim = dim.div(atom_d);
            pending_div = false;
        } else {
            factor *= atom_f;
            dim = dim.mul(atom_d);
        }
    }
    Ok((factor, dim))
}

fn parse_atom(
    s: &str,
    table: &HashMap<&str, UnitValue>,
) -> Result<(f64, Dimension, usize), ParseError> {
    let bytes = s.as_bytes();
    if bytes.is_empty() {
        return Err(ParseError::ValidationError(
            "trailing operator in unit".into(),
        ));
    }
    if bytes[0] == b'(' {
        let mut depth = 0;
        for (j, &b) in bytes.iter().enumerate() {
            if b == b'(' {
                depth += 1;
            } else if b == b')' {
                depth -= 1;
                if depth == 0 {
                    let inner = &s[1..j];
                    let (f, d) = parse_expr(inner, table)?;
                    let mut end = j + 1;
                    let (f2, d2, extra) = apply_power(f, d, &s[end..])?;
                    end += extra;
                    return Ok((f2, d2, end));
                }
            }
        }
        return Err(ParseError::ValidationError("unbalanced '(' in unit".into()));
    }
    // identifier [ ^ number ]
    let mut j = 0;
    while j < bytes.len()
        && (bytes[j].is_ascii_alphanumeric() || bytes[j] == b'_' || bytes[j] > 127)
    {
        j += 1;
    }
    if j == 0 {
        return Err(ParseError::ValidationError(format!(
            "expected unit name in '{s}'"
        )));
    }
    let name = &s[..j];
    let base = table
        .get(name)
        .ok_or_else(|| ParseError::ValidationError(format!("unknown unit '{name}'")))?;
    let mut end = j;
    let (f, d, extra) = apply_power(base.si_factor, base.dim, &s[end..])?;
    end += extra;
    Ok((f, d, end))
}

fn apply_power(f: f64, d: Dimension, rest: &str) -> Result<(f64, Dimension, usize), ParseError> {
    let bytes = rest.as_bytes();
    if bytes.first() == Some(&b'^') {
        let mut k = 1;
        let neg = bytes.get(1) == Some(&b'-');
        if neg {
            k = 2;
        }
        let start = k;
        while k < bytes.len() && (bytes[k].is_ascii_digit() || bytes[k] == b'.') {
            k += 1;
        }
        if k == start {
            return Err(ParseError::ValidationError(
                "expected exponent after ^".into(),
            ));
        }
        let mut p: f64 = rest[start..k]
            .parse()
            .map_err(|_| ParseError::ValidationError("invalid unit exponent".into()))?;
        if neg {
            p = -p;
        }
        return Ok((f.powf(p), d.pow(p), k));
    }
    Ok((f, d, 0))
}

/// Multiplicative factor: `value_in_to = factor * value_in_from`.
pub fn unit_conversion_factor(from_unit: &str, to_unit: &str) -> Result<f64, ParseError> {
    let (f_from, d_from) = parse_unit_expression(from_unit)?;
    let (f_to, d_to) = parse_unit_expression(to_unit)?;
    if !d_from.compatible(d_to) {
        return Err(ParseError::ValidationError(format!(
            "incompatible units '{from_unit}' and '{to_unit}'"
        )));
    }
    Ok(f_from / f_to)
}

/// Validate that `unit` has dimensions appropriate for `quantity`
/// (`length`, `energy`, `mass`, `time`, `velocity`, `force`).
pub fn validate_unit_for_quantity(quantity: &str, unit: &str) -> Result<(), ParseError> {
    let (_, dim) = parse_unit_expression(unit)?;
    let expect = match quantity {
        "length" => Dimension::LENGTH,
        "time" => Dimension::TIME,
        "mass" => Dimension::MASS,
        "energy" => Dimension::MASS
            .mul(Dimension::LENGTH.pow(2.0))
            .div(Dimension::TIME.pow(2.0)),
        "velocity" => Dimension::LENGTH.div(Dimension::TIME),
        "force" => Dimension::MASS
            .mul(Dimension::LENGTH)
            .div(Dimension::TIME.pow(2.0)),
        _ => {
            return Err(ParseError::ValidationError(format!(
                "unknown quantity '{quantity}'"
            )));
        }
    };
    if !dim.compatible(expect) {
        return Err(ParseError::ValidationError(format!(
            "unit '{unit}' is not valid for quantity '{quantity}'"
        )));
    }
    Ok(())
}

/// CON v3 requires `units` object with non-empty `length` and `energy` strings.
pub fn validate_v3_units_metadata(units: &serde_json::Value) -> Result<(), ParseError> {
    let obj = units
        .as_object()
        .ok_or_else(|| ParseError::ValidationError("units must be a JSON object".into()))?;
    for key in ["length", "energy"] {
        let Some(v) = obj.get(key) else {
            return Err(ParseError::ValidationError(format!(
                "v3 units must include non-empty '{key}'"
            )));
        };
        let Some(s) = v.as_str() else {
            return Err(ParseError::ValidationError(format!(
                "units.{key} must be a string"
            )));
        };
        if s.trim().is_empty() {
            return Err(ParseError::ValidationError(format!(
                "units.{key} must be non-empty"
            )));
        }
        validate_unit_for_quantity(key, s)?;
    }
    // Optional keys if present must be valid for their quantity
    for (key, qty) in [
        ("mass", "mass"),
        ("time", "time"),
        ("velocity", "velocity"),
        ("force", "force"),
    ] {
        if let Some(v) = obj.get(key) {
            let s = v.as_str().ok_or_else(|| {
                ParseError::ValidationError(format!("units.{key} must be a string"))
            })?;
            validate_unit_for_quantity(qty, s)?;
        }
    }
    Ok(())
}

/// Default LODE units object for new v3 frames.
pub fn default_v3_units_json() -> serde_json::Value {
    serde_json::json!({
        "length": "angstrom",
        "energy": "eV",
        "mass": "amu",
        "time": "fs"
    })
}

fn preferred_atom(alias: &str) -> Option<&'static str> {
    Some(match alias.trim().to_ascii_lowercase().as_str() {
        "a" | "å" | "angstrom" | "ångstrom" => "angstrom",
        "nm" | "nanometer" | "nanometre" => "nm",
        "m" | "meter" | "metre" => "m",
        "bohr" | "a0" => "bohr",
        "fs" | "femtosecond" | "femtoseconds" => "fs",
        "ps" | "picosecond" | "picoseconds" => "ps",
        "ns" | "nanosecond" | "nanoseconds" => "ns",
        "s" | "sec" | "second" | "seconds" => "s",
        "ev" => "eV",
        "mev" => "meV",
        "hartree" | "ha" => "hartree",
        "j" | "joule" => "J",
        "kj" => "kJ",
        "kcal" => "kcal",
        "amu" | "u" | "dalton" | "da" => "amu",
        "kg" | "kilogram" => "kg",
        "mol" => "mol",
        _ => return None,
    })
}

/// Rewrite a unit expression to preferred names (`A` → `angstrom`, `ev` → `eV`).
/// The expression MUST parse (`unit_conversion_factor` identity).
pub fn canonicalize_unit_expression(expr: &str) -> Result<String, ParseError> {
    let trimmed = expr.trim();
    if trimmed.is_empty() {
        return Err(ParseError::ValidationError("empty unit".into()));
    }
    let _ = unit_conversion_factor(trimmed, trimmed)?;
    let compact: String = trimmed.chars().filter(|c| !c.is_whitespace()).collect();
    let mut out = String::new();
    let bytes = compact.as_bytes();
    let mut i = 0;
    while i < bytes.len() {
        let b = bytes[i];
        if matches!(b, b'*' | b'/' | b'(' | b')' | b'^') {
            if !out.is_empty() && !out.ends_with(' ') && b != b')' && b != b'^' {
                out.push(' ');
            }
            out.push(b as char);
            if b != b'(' && b != b'^' {
                out.push(' ');
            }
            i += 1;
            continue;
        }
        if b == b'-' || b.is_ascii_digit() || b == b'.' {
            out.push(b as char);
            i += 1;
            continue;
        }
        let start = i;
        while i < bytes.len() && (bytes[i].is_ascii_alphabetic() || bytes[i] > 127) {
            i += 1;
        }
        if start == i {
            return Err(ParseError::ValidationError(format!(
                "bad unit token in '{expr}'"
            )));
        }
        let atom = &compact[start..i];
        out.push_str(preferred_atom(atom).unwrap_or(atom));
    }
    Ok(out.split_whitespace().collect::<Vec<_>>().join(" "))
}

/// Rewrite every string member. Does not require v3 `length`+`energy`.
pub fn canonicalize_units_object(
    units: &serde_json::Value,
) -> Result<serde_json::Value, ParseError> {
    let obj = units.as_object().ok_or_else(|| {
        ParseError::ValidationError("units must be a JSON object".into())
    })?;
    let mut out = serde_json::Map::new();
    for (k, v) in obj {
        let s = v.as_str().ok_or_else(|| {
            ParseError::ValidationError(format!("units.{k} must be a string"))
        })?;
        out.insert(
            k.clone(),
            serde_json::Value::String(canonicalize_unit_expression(s)?),
        );
    }
    Ok(serde_json::Value::Object(out))
}

/// v3 validate then canonicalize (preferred names on line 2).
pub fn canonicalize_units_metadata(
    units: &serde_json::Value,
) -> Result<serde_json::Value, ParseError> {
    validate_v3_units_metadata(units)?;
    canonicalize_units_object(units)
}

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

    #[test]
    fn angstrom_to_nm() {
        let f = unit_conversion_factor("angstrom", "nm").unwrap();
        assert!((f - 0.1).abs() < 1e-12);
    }

    #[test]
    fn ev_to_mev() {
        let f = unit_conversion_factor("eV", "meV").unwrap();
        assert!((f - 1000.0).abs() < 1e-6);
    }

    #[test]
    fn dimensional_mismatch() {
        assert!(unit_conversion_factor("angstrom", "eV").is_err());
    }

    #[test]
    fn ev_to_kcal_per_mol_is_metatomic() {
        let f = unit_conversion_factor("eV", "kcal / mol").unwrap();
        assert!((f - 23.060_548).abs() < 1e-3, "got {f}");
        let force = unit_conversion_factor("eV / angstrom", "kJ / mol / angstrom").unwrap();
        assert!((force - 96.485_332).abs() < 1e-3, "got {force}");
        let ns = unit_conversion_factor("ns", "ps").unwrap();
        assert!((ns - 1000.0).abs() < 1e-9, "got {ns}");
    }

    #[test]
    fn validate_length_energy() {
        validate_unit_for_quantity("length", "nm").unwrap();
        validate_unit_for_quantity("energy", "hartree").unwrap();
        assert!(validate_unit_for_quantity("energy", "angstrom").is_err());
    }

    #[test]
    fn v3_units_require_length_energy() {
        assert!(validate_v3_units_metadata(&serde_json::json!({"length": "A"})).is_err());
        validate_v3_units_metadata(&default_v3_units_json()).unwrap();
    }

    #[test]
    fn canonicalize_aliases() {
        assert_eq!(canonicalize_unit_expression("A").unwrap(), "angstrom");
        assert_eq!(canonicalize_unit_expression("ev").unwrap(), "eV");
        assert_eq!(canonicalize_unit_expression("femtosecond").unwrap(), "fs");
        assert_eq!(
            canonicalize_unit_expression("eV/angstrom").unwrap(),
            "eV / angstrom"
        );
        let u = canonicalize_units_metadata(&serde_json::json!({
            "length": "A",
            "energy": "ev",
            "time": "femtosecond"
        }))
        .unwrap();
        assert_eq!(u["length"], "angstrom");
        assert_eq!(u["energy"], "eV");
        assert_eq!(u["time"], "fs");
    }
}