oxiphysics-io 0.1.1

// Copyright 2026 COOLJAPAN OU (Team KitaSan)
// SPDX-License-Identifier: Apache-2.0

//! Scientific data format support for OxiPhysics.
//!
//! Provides read/write support for Gaussian Cube files, Molden geometry files,
//! extended XYZ (extxyz) trajectory frames, CIF cell parameter lines, and a
//! simplified parser for Gaussian formatted checkpoint (fchk) arrays.

use std::collections::HashMap;
use std::io::{BufRead, Write};

// ── Cube file ─────────────────────────────────────────────────────────────────

/// A Gaussian Cube volumetric data file.
#[derive(Debug, Clone)]
pub struct CubeFile {
    /// Number of atoms in the molecule.
    pub n_atoms: usize,
    /// Origin of the volumetric grid in Bohr.
    pub origin: [f64; 3],
    /// Axis vectors `axes[i]` define the step along axis `i`.
    pub axes: [[f64; 3]; 3],
    /// Number of grid points along the X axis.
    pub nx: usize,
    /// Number of grid points along the Y axis.
    pub ny: usize,
    /// Number of grid points along the Z axis.
    pub nz: usize,
    /// Volumetric data in row-major order `[x][y][z]`.
    pub data: Vec<f64>,
    /// Atomic positions in Bohr, one per atom.
    pub atom_positions: Vec<[f64; 3]>,
    /// Atomic numbers, one per atom.
    pub atom_numbers: Vec<u32>,
}

/// Write a [`CubeFile`] to `path` in Gaussian Cube format.
pub fn write_cube(path: &str, cf: &CubeFile) -> Result<(), std::io::Error> {
    let file = std::fs::File::create(path)?;
    let mut w = std::io::BufWriter::new(file);

    writeln!(w, " Cube file generated by OxiPhysics")?;
    writeln!(w, " Volumetric data")?;
    writeln!(
        w,
        " {:5} {:12.6} {:12.6} {:12.6}",
        cf.n_atoms as i64, cf.origin[0], cf.origin[1], cf.origin[2]
    )?;
    writeln!(
        w,
        " {:5} {:12.6} {:12.6} {:12.6}",
        cf.nx, cf.axes[0][0], cf.axes[0][1], cf.axes[0][2]
    )?;
    writeln!(
        w,
        " {:5} {:12.6} {:12.6} {:12.6}",
        cf.ny, cf.axes[1][0], cf.axes[1][1], cf.axes[1][2]
    )?;
    writeln!(
        w,
        " {:5} {:12.6} {:12.6} {:12.6}",
        cf.nz, cf.axes[2][0], cf.axes[2][1], cf.axes[2][2]
    )?;
    for i in 0..cf.n_atoms {
        let z = cf.atom_numbers.get(i).copied().unwrap_or(0);
        let pos = cf.atom_positions.get(i).copied().unwrap_or([0.0; 3]);
        writeln!(
            w,
            " {:5} {:12.6} {:12.6} {:12.6} {:12.6}",
            z, z as f64, pos[0], pos[1], pos[2]
        )?;
    }
    // Write volumetric data, 6 values per line
    for (idx, &val) in cf.data.iter().enumerate() {
        if idx > 0 && idx % 6 == 0 {
            writeln!(w)?;
        }
        write!(w, " {:12.5E}", val)?;
    }
    writeln!(w)?;
    Ok(())
}

/// Read a [`CubeFile`] from `path`.
pub fn read_cube(path: &str) -> Result<CubeFile, std::io::Error> {
    let file = std::fs::File::open(path)?;
    let reader = std::io::BufReader::new(file);
    let mut lines = reader.lines();

    // Skip two comment lines
    lines.next();
    lines.next();

    let parse_err = |s: &str| std::io::Error::new(std::io::ErrorKind::InvalidData, s.to_string());

    let read_line = |lines: &mut dyn Iterator<Item = Result<String, std::io::Error>>| {
        lines
            .next()
            .ok_or_else(|| std::io::Error::new(std::io::ErrorKind::UnexpectedEof, "EOF"))
            .and_then(|r| r)
    };

    // Line: n_atoms  ox oy oz
    let atom_line = read_line(&mut lines)?;
    let toks: Vec<&str> = atom_line.split_whitespace().collect();
    let n_atoms = toks
        .first()
        .ok_or_else(|| parse_err("missing n_atoms"))?
        .parse::<i64>()
        .map_err(|e| parse_err(&e.to_string()))?
        .unsigned_abs() as usize;
    let ox: f64 = toks.get(1).unwrap_or(&"0").parse().unwrap_or(0.0);
    let oy: f64 = toks.get(2).unwrap_or(&"0").parse().unwrap_or(0.0);
    let oz: f64 = toks.get(3).unwrap_or(&"0").parse().unwrap_or(0.0);

    let mut axes = [[0.0_f64; 3]; 3];
    let mut dims = [0usize; 3];
    for i in 0..3 {
        let l = read_line(&mut lines)?;
        let t: Vec<&str> = l.split_whitespace().collect();
        dims[i] = t.first().unwrap_or(&"0").parse().unwrap_or(0);
        axes[i][0] = t.get(1).unwrap_or(&"0").parse().unwrap_or(0.0);
        axes[i][1] = t.get(2).unwrap_or(&"0").parse().unwrap_or(0.0);
        axes[i][2] = t.get(3).unwrap_or(&"0").parse().unwrap_or(0.0);
    }

    let mut atom_numbers = Vec::with_capacity(n_atoms);
    let mut atom_positions = Vec::with_capacity(n_atoms);
    for _ in 0..n_atoms {
        let l = read_line(&mut lines)?;
        let t: Vec<&str> = l.split_whitespace().collect();
        let z: u32 = t.first().unwrap_or(&"0").parse().unwrap_or(0);
        let x: f64 = t.get(2).unwrap_or(&"0").parse().unwrap_or(0.0);
        let y: f64 = t.get(3).unwrap_or(&"0").parse().unwrap_or(0.0);
        let zc: f64 = t.get(4).unwrap_or(&"0").parse().unwrap_or(0.0);
        atom_numbers.push(z);
        atom_positions.push([x, y, zc]);
    }

    let total = dims[0] * dims[1] * dims[2];
    let mut data = Vec::with_capacity(total);
    for line in lines {
        let l = line?;
        for tok in l.split_whitespace() {
            if let Ok(v) = tok.parse::<f64>() {
                data.push(v);
            }
        }
    }

    Ok(CubeFile {
        n_atoms,
        origin: [ox, oy, oz],
        axes,
        nx: dims[0],
        ny: dims[1],
        nz: dims[2],
        data,
        atom_positions,
        atom_numbers,
    })
}

// ── Molden file ───────────────────────────────────────────────────────────────

/// A simplified Molden file containing atomic geometry and vibrational data.
#[derive(Debug, Clone)]
pub struct MoldenFile {
    /// Atoms as `(element_symbol, [x, y, z])` in Ångström.
    pub atoms: Vec<(String, [f64; 3])>,
    /// Vibrational frequencies in cm⁻¹.
    pub frequencies: Vec<f64>,
    /// IR intensities in km/mol.
    pub intensities: Vec<f64>,
}

/// Write Molden geometry section to `path`.
pub fn write_molden_geometry(path: &str, mf: &MoldenFile) -> Result<(), std::io::Error> {
    let file = std::fs::File::create(path)?;
    let mut w = std::io::BufWriter::new(file);
    writeln!(w, "[Molden Format]")?;
    writeln!(w, "[Atoms] Angs")?;
    for (i, (sym, pos)) in mf.atoms.iter().enumerate() {
        writeln!(
            w,
            " {:<4} {:5} {:3} {:14.8} {:14.8} {:14.8}",
            sym,
            i + 1,
            1, // dummy atomic number
            pos[0],
            pos[1],
            pos[2]
        )?;
    }
    if !mf.frequencies.is_empty() {
        writeln!(w, "[FREQ]")?;
        for &f in &mf.frequencies {
            writeln!(w, " {:14.6}", f)?;
        }
    }
    if !mf.intensities.is_empty() {
        writeln!(w, "[INT]")?;
        for &ir in &mf.intensities {
            writeln!(w, " {:14.6}", ir)?;
        }
    }
    Ok(())
}

// ── Extended XYZ ──────────────────────────────────────────────────────────────

/// A single frame from an extended XYZ trajectory.
#[derive(Debug, Clone)]
pub struct XyzFrame {
    /// Simulation step index.
    pub step: usize,
    /// Simulation time in ps.
    pub time: f64,
    /// Atoms as `(element_symbol, [x, y, z])`.
    pub atoms: Vec<(String, [f64; 3])>,
}

/// Parse key=value pairs from an extXYZ comment line.
///
/// Handles quoted string values and bare numeric values.
pub fn extxyz_parse_comment(line: &str) -> HashMap<String, String> {
    let mut map = HashMap::new();
    // Simple tokeniser: split on spaces that are not inside quotes
    let mut chars = line.chars().peekable();
    loop {
        // Skip whitespace
        while chars.peek() == Some(&' ') {
            chars.next();
        }
        if chars.peek().is_none() {
            break;
        }
        // Read key
        let key: String = chars.by_ref().take_while(|&c| c != '=').collect();
        let key = key.trim().to_string();
        if key.is_empty() {
            break;
        }
        // Read value
        let value = if chars.peek() == Some(&'"') {
            chars.next(); // skip opening quote
            let v: String = chars.by_ref().take_while(|&c| c != '"').collect();
            v
        } else {
            // Read until space
            let v: String = chars.by_ref().take_while(|&c| c != ' ').collect();
            v
        };
        if !key.is_empty() {
            map.insert(key, value);
        }
    }
    map
}

/// Write an extXYZ frame to `path`.
pub fn write_extxyz_frame(
    path: &str,
    frame: &XyzFrame,
    properties: &HashMap<String, String>,
) -> Result<(), std::io::Error> {
    let file = std::fs::File::create(path)?;
    let mut w = std::io::BufWriter::new(file);
    writeln!(w, "{}", frame.atoms.len())?;
    // Comment line with step, time, and extra properties
    let mut comment = format!("step={} time={:.6}", frame.step, frame.time);
    for (k, v) in properties {
        comment.push_str(&format!(" {}={}", k, v));
    }
    writeln!(w, "{}", comment)?;
    for (sym, pos) in &frame.atoms {
        writeln!(w, "{} {:14.8} {:14.8} {:14.8}", sym, pos[0], pos[1], pos[2])?;
    }
    Ok(())
}

// ── CIF helpers ───────────────────────────────────────────────────────────────

/// Format a CIF `_cell_*` section for the given unit cell parameters.
///
/// Angles are in degrees; lengths in Ångström.
pub fn cif_cell_parameters_line(
    a: f64,
    b: f64,
    c: f64,
    alpha: f64,
    beta: f64,
    gamma: f64,
) -> String {
    format!(
        "_cell_length_a                  {:.4}\n\
         _cell_length_b                  {:.4}\n\
         _cell_length_c                  {:.4}\n\
         _cell_angle_alpha               {:.4}\n\
         _cell_angle_beta                {:.4}\n\
         _cell_angle_gamma               {:.4}",
        a, b, c, alpha, beta, gamma
    )
}

// ── Gaussian fchk parser ──────────────────────────────────────────────────────

/// Read a float-valued array from a Gaussian formatted checkpoint file.
///
/// Searches for a section matching `keyword` and reads `N` values where `N`
/// is declared on the keyword line.  Returns an empty `Vec` if the keyword is
/// not found.
pub fn fchk_read_array(path: &str, keyword: &str) -> Result<Vec<f64>, std::io::Error> {
    let file = std::fs::File::open(path)?;
    let reader = std::io::BufReader::new(file);
    let lines: Vec<String> = reader.lines().collect::<Result<_, _>>()?;

    let mut reading = false;
    let mut count = 0usize;
    let mut values: Vec<f64> = Vec::new();

    for line in &lines {
        if !reading {
            if line.contains(keyword) && line.contains("N=") {
                // Extract count after "N="
                if let Some(pos) = line.find("N=") {
                    let rest = line[pos + 2..].trim();
                    if let Ok(n) = rest.parse::<usize>() {
                        count = n;
                        reading = true;
                        values.reserve(count);
                    }
                }
            }
        } else {
            for tok in line.split_whitespace() {
                if let Ok(v) = tok.parse::<f64>() {
                    values.push(v);
                    if values.len() >= count {
                        return Ok(values);
                    }
                }
            }
        }
    }
    Ok(values)
}

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

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

    fn sample_cube() -> CubeFile {
        CubeFile {
            n_atoms: 2,
            origin: [0.0, 0.0, 0.0],
            axes: [[0.2, 0.0, 0.0], [0.0, 0.2, 0.0], [0.0, 0.0, 0.2]],
            nx: 2,
            ny: 2,
            nz: 2,
            data: vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0],
            atom_positions: vec![[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]],
            atom_numbers: vec![6, 1],
        }
    }

    #[test]
    fn test_write_read_cube_roundtrip() {
        let path = "/tmp/test_oxiphysics_cube.cube";
        let cf = sample_cube();
        write_cube(path, &cf).unwrap();
        let cf2 = read_cube(path).unwrap();
        assert_eq!(cf2.n_atoms, 2);
        assert_eq!(cf2.nx, 2);
        assert_eq!(cf2.ny, 2);
        assert_eq!(cf2.nz, 2);
        assert_eq!(cf2.data.len(), 8);
    }

    #[test]
    fn test_cube_data_values() {
        let path = "/tmp/test_oxiphysics_cube_data.cube";
        let cf = sample_cube();
        write_cube(path, &cf).unwrap();
        let cf2 = read_cube(path).unwrap();
        assert!((cf2.data[0] - 1.0).abs() < 1e-4);
        assert!((cf2.data[7] - 8.0).abs() < 1e-4);
    }

    #[test]
    fn test_cube_atom_numbers() {
        let path = "/tmp/test_oxiphysics_cube_atoms.cube";
        let cf = sample_cube();
        write_cube(path, &cf).unwrap();
        let cf2 = read_cube(path).unwrap();
        assert_eq!(cf2.atom_numbers[0], 6);
        assert_eq!(cf2.atom_numbers[1], 1);
    }

    #[test]
    fn test_cube_atom_positions() {
        let path = "/tmp/test_oxiphysics_cube_pos.cube";
        let cf = sample_cube();
        write_cube(path, &cf).unwrap();
        let cf2 = read_cube(path).unwrap();
        assert!((cf2.atom_positions[1][0] - 1.0).abs() < 1e-4);
    }

    #[test]
    fn test_cube_origin() {
        let path = "/tmp/test_oxiphysics_cube_origin.cube";
        let cf = sample_cube();
        write_cube(path, &cf).unwrap();
        let cf2 = read_cube(path).unwrap();
        assert!((cf2.origin[0]).abs() < 1e-6);
    }

    #[test]
    fn test_write_molden_creates_file() {
        let path = "/tmp/test_oxiphysics_molden.mld";
        let mf = MoldenFile {
            atoms: vec![
                ("C".to_string(), [0.0, 0.0, 0.0]),
                ("H".to_string(), [1.0, 0.0, 0.0]),
            ],
            frequencies: vec![1000.0, 2000.0],
            intensities: vec![50.0, 100.0],
        };
        write_molden_geometry(path, &mf).unwrap();
        let content = std::fs::read_to_string(path).unwrap();
        assert!(content.contains("[Molden Format]"));
        assert!(content.contains("[Atoms]"));
        assert!(content.contains("[FREQ]"));
    }

    #[test]
    fn test_molden_no_freq() {
        let path = "/tmp/test_oxiphysics_molden_nofreq.mld";
        let mf = MoldenFile {
            atoms: vec![("O".to_string(), [0.0, 0.0, 0.0])],
            frequencies: vec![],
            intensities: vec![],
        };
        write_molden_geometry(path, &mf).unwrap();
        let content = std::fs::read_to_string(path).unwrap();
        assert!(!content.contains("[FREQ]"));
    }

    #[test]
    fn test_extxyz_parse_simple() {
        let line =
            "Lattice=\"5.0 0 0 0 5 0 0 0 5\" Properties=species:S:1:pos:R:3 step=10 time=0.5";
        let map = extxyz_parse_comment(line);
        assert_eq!(map.get("step").map(|s| s.as_str()), Some("10"));
        assert_eq!(map.get("time").map(|s| s.as_str()), Some("0.5"));
    }

    #[test]
    fn test_extxyz_parse_empty() {
        let map = extxyz_parse_comment("");
        assert!(map.is_empty());
    }

    #[test]
    fn test_extxyz_parse_quoted_value() {
        let line = r#"Lattice="1 0 0 0 1 0 0 0 1" step=5"#;
        let map = extxyz_parse_comment(line);
        assert!(map.contains_key("Lattice"));
        assert_eq!(map.get("step").map(|s| s.as_str()), Some("5"));
    }

    #[test]
    fn test_write_extxyz_frame() {
        let path = "/tmp/test_oxiphysics_extxyz.xyz";
        let frame = XyzFrame {
            step: 10,
            time: 0.1,
            atoms: vec![
                ("C".to_string(), [0.0, 0.0, 0.0]),
                ("H".to_string(), [1.1, 0.0, 0.0]),
            ],
        };
        let mut props = HashMap::new();
        props.insert("energy".to_string(), "-100.5".to_string());
        write_extxyz_frame(path, &frame, &props).unwrap();
        let content = std::fs::read_to_string(path).unwrap();
        assert!(content.starts_with("2\n"));
        assert!(content.contains("step=10"));
        assert!(content.contains("energy=-100.5"));
    }

    #[test]
    fn test_write_extxyz_atom_count() {
        let path = "/tmp/test_oxiphysics_extxyz_count.xyz";
        let frame = XyzFrame {
            step: 0,
            time: 0.0,
            atoms: vec![("N".to_string(), [0.0, 0.0, 0.0])],
        };
        write_extxyz_frame(path, &frame, &HashMap::new()).unwrap();
        let content = std::fs::read_to_string(path).unwrap();
        assert!(content.starts_with("1\n"));
    }

    #[test]
    fn test_cif_cell_line_format() {
        let line = cif_cell_parameters_line(5.0, 5.0, 5.0, 90.0, 90.0, 90.0);
        assert!(line.contains("_cell_length_a"));
        assert!(line.contains("5.0000"));
        assert!(line.contains("90.0000"));
    }

    #[test]
    fn test_cif_cell_line_non_cubic() {
        let line = cif_cell_parameters_line(3.0, 4.0, 5.0, 80.0, 100.0, 120.0);
        assert!(line.contains("3.0000"));
        assert!(line.contains("4.0000"));
        assert!(line.contains("120.0000"));
    }

    #[test]
    fn test_fchk_read_array_not_found() {
        let path = "/tmp/test_fchk_empty.fchk";
        std::fs::write(path, "Some data\n").unwrap();
        let arr = fchk_read_array(path, "NonExistentKeyword").unwrap();
        assert!(arr.is_empty());
    }

    #[test]
    fn test_fchk_read_array_found() {
        let path = "/tmp/test_fchk_data.fchk";
        let content = "Total Energy                       R   N=           3\n   1.0000000E+00   2.0000000E+00   3.0000000E+00\n";
        std::fs::write(path, content).unwrap();
        let arr = fchk_read_array(path, "Total Energy").unwrap();
        assert_eq!(arr.len(), 3);
        assert!((arr[0] - 1.0).abs() < 1e-6);
        assert!((arr[2] - 3.0).abs() < 1e-6);
    }

    #[test]
    fn test_cube_file_no_atoms() {
        let path = "/tmp/test_oxiphysics_cube_zero.cube";
        let cf = CubeFile {
            n_atoms: 0,
            origin: [0.0; 3],
            axes: [[0.1, 0.0, 0.0], [0.0, 0.1, 0.0], [0.0, 0.0, 0.1]],
            nx: 1,
            ny: 1,
            nz: 1,
            data: vec![0.5],
            atom_positions: vec![],
            atom_numbers: vec![],
        };
        write_cube(path, &cf).unwrap();
        let cf2 = read_cube(path).unwrap();
        assert_eq!(cf2.n_atoms, 0);
    }

    #[test]
    fn test_molden_atom_symbols() {
        let path = "/tmp/test_oxiphysics_molden_sym.mld";
        let mf = MoldenFile {
            atoms: vec![("Fe".to_string(), [0.5, 0.5, 0.5])],
            frequencies: vec![],
            intensities: vec![],
        };
        write_molden_geometry(path, &mf).unwrap();
        let content = std::fs::read_to_string(path).unwrap();
        assert!(content.contains("Fe"));
    }

    #[test]
    fn test_extxyz_parse_multiple_keys() {
        let line = "a=1 b=2 c=3";
        let map = extxyz_parse_comment(line);
        assert_eq!(map.len(), 3);
        assert_eq!(map.get("a").map(|s| s.as_str()), Some("1"));
    }

    #[test]
    fn test_write_extxyz_time() {
        let path = "/tmp/test_oxiphysics_extxyz_time.xyz";
        let frame = XyzFrame {
            step: 5,
            time: 1.23456,
            atoms: vec![("H".to_string(), [0.0, 0.0, 0.0])],
        };
        write_extxyz_frame(path, &frame, &HashMap::new()).unwrap();
        let content = std::fs::read_to_string(path).unwrap();
        assert!(content.contains("time=1.234560"));
    }

    #[test]
    fn test_cif_cell_contains_all_fields() {
        let line = cif_cell_parameters_line(1.0, 2.0, 3.0, 91.0, 92.0, 93.0);
        for field in &[
            "_cell_length_a",
            "_cell_length_b",
            "_cell_length_c",
            "_cell_angle_alpha",
            "_cell_angle_beta",
            "_cell_angle_gamma",
        ] {
            assert!(line.contains(field), "Missing field: {}", field);
        }
    }
}