readcon-db 0.2.0

//! Canonical unit spellings (same preferred-atom table as
//! `readcon_core::units::canonicalize_unit_expression` on core git main).
//! crates.io `readcon-core` 0.14 has the SI parser but not that symbol,
//! so this crate keeps the rewrite here. Callers may write `A`, `ev`,
//! `femtosecond`; metadata stores `angstrom`, `eV`, `fs`.

use crate::error::{Error, Result};

/// Natural time unit eOn and ASE write: sqrt(amu * Angstrom^2 / eV).
///
/// readcon-core 0.16.2 parses this as a time (about 10.180505710759414 fs).
/// eOn hard-codes 10.1805055 fs. D10.
///
/// R. Goswami, MethodsX 2026, doi:10.1016/j.mex.2026.103899.
/// R. Goswami, MethodsX 2026, doi:10.1016/j.mex.2026.103851.
/// R. Goswami et al., J. Chem. Theory Comput. 2025, doi:10.1021/acs.jctc.5c00866.
/// S. T. Chill et al., Modelling and Simulation in Materials Science and Engineering 22 (2014) 055002, doi:10.1088/0965-0393/22/5/055002.
pub const NATURAL_TIME_UNIT: &str = "(amu*angstrom^2/eV)^0.5";

/// Minor from line 2. Absent means v3.0 when `spec_version` is 3 (D10, D15).
pub fn con_spec_minor(header: &readcon_core::types::FrameHeader) -> Option<u64> {
    header
        .metadata
        .get("con_spec_minor")
        .and_then(|v| v.as_u64())
}

/// From v3.1 on, `units.time` is the velocity time unit.
/// v1, v2, and v3.0 (no `con_spec_minor`) keep the natural unit whatever the label says.
pub fn velocity_time_is_declared(header: &readcon_core::types::FrameHeader) -> bool {
    if header.spec_version > 3 {
        return true;
    }
    header.spec_version == 3 && con_spec_minor(header).unwrap_or(0) >= 1
}

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

/// Validate with the SI parser and rewrite aliases to preferred names.
pub fn canonicalize_unit(expr: &str) -> Result<String> {
    let trimmed = expr.trim();
    if trimmed.is_empty() {
        return Err(Error::Message("empty unit".into()));
    }
    readcon_core::units::unit_conversion_factor(trimmed, trimmed)
        .map_err(|e| Error::Message(e.to_string()))?;
    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(Error::Message(format!("bad unit token in '{expr}'")));
        }
        let atom = &compact[start..i];
        let pref = preferred_atom(atom).unwrap_or(atom);
        out.push_str(pref);
    }
    Ok(out.split_whitespace().collect::<Vec<_>>().join(" "))
}

pub fn canonicalize_units_object(units: serde_json::Value) -> Result<serde_json::Value> {
    let obj = units
        .as_object()
        .ok_or_else(|| Error::Message("units must be an object".into()))?;
    let mut out = serde_json::Map::new();
    for (k, v) in obj {
        let Some(s) = v.as_str() else {
            return Err(Error::Message(format!("units.{k} must be a string")));
        };
        out.insert(k.clone(), serde_json::Value::String(canonicalize_unit(s)?));
    }
    Ok(serde_json::Value::Object(out))
}

/// LODE CON defaults when a dimension is missing (spec v3).
pub fn default_con_units() -> serde_json::Value {
    serde_json::json!({
        "length": "angstrom",
        "energy": "eV",
        "mass": "amu",
        "time": "fs"
    })
}

fn dim_unit<'a>(u: &'a serde_json::Value, dim: &str, default: &'a str) -> &'a str {
    u.get(dim).and_then(|v| v.as_str()).unwrap_or(default)
}

fn dim_factor(
    old: &serde_json::Value,
    new: &serde_json::Value,
    dim: &str,
    default: &str,
) -> Result<f64> {
    let from = dim_unit(old, dim, default);
    let to = dim_unit(new, dim, from);
    if from == to {
        return Ok(1.0);
    }
    readcon_core::units::unit_conversion_factor(from, to).map_err(|e| Error::Message(e.to_string()))
}

/// Merge `new` onto existing (or CON default) units. Rewrites aliases.
pub fn merge_units(
    old: Option<&serde_json::Value>,
    new: serde_json::Value,
) -> Result<serde_json::Value> {
    let new = canonicalize_units_object(new)?;
    let mut out = match old.and_then(|v| v.as_object()) {
        Some(m) => m.clone(),
        None => default_con_units().as_object().cloned().unwrap_or_default(),
    };
    if let Some(n) = new.as_object() {
        for (k, v) in n {
            out.insert(k.clone(), v.clone());
        }
    }
    Ok(serde_json::Value::Object(out))
}

/// Scale stored numbers so `new_units` is an honest label (not a relabel).
pub fn rescale_frame_units(
    fr: &mut readcon_core::types::ConFrame,
    new_units: &serde_json::Value,
) -> Result<()> {
    let old = fr.header.units().cloned().unwrap_or_else(default_con_units);
    let lf = dim_factor(&old, new_units, "length", "angstrom")?;
    let ef = dim_factor(&old, new_units, "energy", "eV")?;
    let tf = dim_factor(&old, new_units, "time", "fs")?;
    let mf = dim_factor(&old, new_units, "mass", "amu")?;
    // D10: pre-v3.1 velocities stay in the natural unit. The time label
    // still rescales the clock (`time`, `timestep`) and is not a velocity unit.
    let vel_tf = if velocity_time_is_declared(&fr.header) {
        tf
    } else {
        1.0
    };
    let vf = lf / vel_tf;
    let ff = ef / lf;
    if lf != 1.0 {
        fr.header.boxl[0] *= lf;
        fr.header.boxl[1] *= lf;
        fr.header.boxl[2] *= lf;
        if let Some(arr) = fr.header.metadata.get_mut("lattice_vectors") {
            if let Some(rows) = arr.as_array_mut() {
                for row in rows {
                    if let Some(cs) = row.as_array_mut() {
                        for c in cs {
                            if let Some(x) = c.as_f64() {
                                *c = serde_json::Value::from(x * lf);
                            }
                        }
                    }
                }
            }
        }
    }
    if tf != 1.0 {
        if let Some(t) = fr.header.time() {
            fr.header.set_time(t * tf);
        }
        if let Some(dt) = fr.header.timestep() {
            fr.header.set_timestep(dt * tf);
        }
    }
    if ef != 1.0 {
        if let Some(e) = fr.header.energy() {
            fr.header.set_energy(e * ef);
        }
    }
    let n = fr.positions.nrows();
    for i in 0..n {
        if lf != 1.0 {
            let mut p = fr.positions.as_f64_row(i);
            p[0] *= lf;
            p[1] *= lf;
            p[2] *= lf;
            fr.positions.set_f64_row(i, p);
        }
        if fr.velocities.nrows() == n && vf != 1.0 {
            let mut v = fr.velocities.as_f64_row(i);
            v[0] *= vf;
            v[1] *= vf;
            v[2] *= vf;
            fr.velocities.set_f64_row(i, v);
        }
        if fr.forces.nrows() == n && ff != 1.0 {
            let mut f = fr.forces.as_f64_row(i);
            f[0] *= ff;
            f[1] *= ff;
            f[2] *= ff;
            fr.forces.set_f64_row(i, f);
        }
        if fr.atom_energies.len() == n && ef != 1.0 {
            let e = fr.atom_energies.get_f64(i);
            fr.atom_energies.set_f64(i, e * ef);
        }
        if fr.masses.len() == n && mf != 1.0 {
            let m = fr.masses.get_f64(i);
            fr.masses.set_f64(i, m * mf);
        }
        if let Some(a) = fr.atom_data.get_mut(i) {
            if lf != 1.0 {
                a.x *= lf;
                a.y *= lf;
                a.z *= lf;
            }
            if let Some(mut v) = a.velocity() {
                if vf != 1.0 {
                    v[0] *= vf;
                    v[1] *= vf;
                    v[2] *= vf;
                    a.set_velocity(Some(v));
                }
            }
            if let Some(mut f) = a.force() {
                if ff != 1.0 {
                    f[0] *= ff;
                    f[1] *= ff;
                    f[2] *= ff;
                    a.set_force(Some(f));
                }
            }
            if let Some(e) = a.energy() {
                if ef != 1.0 {
                    a.set_energy(Some(e * ef));
                }
            }
        }
    }
    if mf != 1.0 {
        for m in &mut fr.header.masses_per_type {
            *m *= mf;
        }
    }
    Ok(())
}

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

    #[test]
    fn aliases_to_preferred() {
        assert_eq!(canonicalize_unit("A").unwrap(), "angstrom");
        assert_eq!(canonicalize_unit("ev").unwrap(), "eV");
        assert_eq!(canonicalize_unit("femtosecond").unwrap(), "fs");
        assert_eq!(canonicalize_unit("eV/angstrom").unwrap(), "eV / angstrom");
        let u = canonicalize_units_object(serde_json::json!({"time": "femtosecond"})).unwrap();
        assert_eq!(u["time"], "fs");
        assert!(u.get("length").is_none());
    }
}