vtk-pure-rs 0.2.0

use std::io::Write;
use std::path::Path;

use crate::data::PolyData;
use crate::types::VtkError;

/// Writer for EnSight Gold format (ASCII).
///
/// Produces a `.case` file, a `.geo` geometry file, and optional `.scl`/`.vec`
/// variable files. This is a widely-supported format readable by ParaView,
/// EnSight, and many other visualization tools.
pub struct EnSightWriter;

impl EnSightWriter {
    /// Write a PolyData mesh to EnSight Gold format.
    ///
    /// Creates `{base_name}.case`, `{base_name}.geo`, and optional variable files
    /// in the given directory.
    pub fn write(dir: &Path, base_name: &str, poly_data: &PolyData) -> Result<(), VtkError> {
        std::fs::create_dir_all(dir)?;

        let geo_name = format!("{base_name}.geo");
        let case_path = dir.join(format!("{base_name}.case"));
        let geo_path = dir.join(&geo_name);

        // Collect scalar and vector variable names
        let mut scalar_names = Vec::new();
        let mut vector_names = Vec::new();
        for i in 0..poly_data.point_data().num_arrays() {
            if let Some(arr) = poly_data.point_data().get_array_by_index(i) {
                let name = arr.name().to_string();
                match arr.num_components() {
                    1 => scalar_names.push(name),
                    3 => vector_names.push(name),
                    _ => {}
                }
            }
        }

        // Write case file
        write_case(&case_path, &geo_name, base_name, &scalar_names, &vector_names)?;

        // Write geometry file
        write_geometry(&geo_path, poly_data)?;

        // Write variable files
        for name in &scalar_names {
            let var_path = dir.join(format!("{base_name}.{name}.scl"));
            write_scalar_variable(&var_path, poly_data, name)?;
        }
        for name in &vector_names {
            let var_path = dir.join(format!("{base_name}.{name}.vec"));
            write_vector_variable(&var_path, poly_data, name)?;
        }

        Ok(())
    }
}

fn write_case(
    path: &Path,
    geo_name: &str,
    base_name: &str,
    scalar_names: &[String],
    vector_names: &[String],
) -> Result<(), VtkError> {
    let mut f = std::fs::File::create(path)?;
    writeln!(f, "FORMAT")?;
    writeln!(f, "type: ensight gold")?;
    writeln!(f)?;
    writeln!(f, "GEOMETRY")?;
    writeln!(f, "model: {geo_name}")?;

    if !scalar_names.is_empty() || !vector_names.is_empty() {
        writeln!(f)?;
        writeln!(f, "VARIABLE")?;
        for name in scalar_names {
            writeln!(f, "scalar per node: {name} {base_name}.{name}.scl")?;
        }
        for name in vector_names {
            writeln!(f, "vector per node: {name} {base_name}.{name}.vec")?;
        }
    }

    Ok(())
}

fn write_geometry(path: &Path, pd: &PolyData) -> Result<(), VtkError> {
    let mut f = std::fs::File::create(path)?;

    writeln!(f, "EnSight Gold geometry file")?;
    writeln!(f, "Generated by vtk-rs")?;
    writeln!(f, "node id off")?;
    writeln!(f, "element id off")?;
    writeln!(f, "part")?;
    writeln!(f, "{:>10}", 1)?;
    writeln!(f, "mesh")?;
    writeln!(f, "coordinates")?;
    writeln!(f, "{:>10}", pd.points.len())?;

    // X coordinates
    for i in 0..pd.points.len() {
        let p = pd.points.get(i);
        writeln!(f, "{:>12.5e}", p[0] as f32)?;
    }
    // Y coordinates
    for i in 0..pd.points.len() {
        let p = pd.points.get(i);
        writeln!(f, "{:>12.5e}", p[1] as f32)?;
    }
    // Z coordinates
    for i in 0..pd.points.len() {
        let p = pd.points.get(i);
        writeln!(f, "{:>12.5e}", p[2] as f32)?;
    }

    // Triangles
    let mut tri_cells = Vec::new();
    for cell in pd.polys.iter() {
        if cell.len() < 3 {
            continue;
        }
        for i in 1..cell.len() - 1 {
            tri_cells.push([cell[0] as usize, cell[i] as usize, cell[i + 1] as usize]);
        }
    }

    if !tri_cells.is_empty() {
        writeln!(f, "tria3")?;
        writeln!(f, "{:>10}", tri_cells.len())?;
        for tri in &tri_cells {
            // EnSight uses 1-based indices
            writeln!(f, "{:>10}{:>10}{:>10}", tri[0] + 1, tri[1] + 1, tri[2] + 1)?;
        }
    }

    Ok(())
}

fn write_scalar_variable(path: &Path, pd: &PolyData, name: &str) -> Result<(), VtkError> {
    let arr = pd.point_data().get_array(name)
        .ok_or_else(|| VtkError::InvalidData(format!("array '{name}' not found")))?;

    let mut f = std::fs::File::create(path)?;
    writeln!(f, "{name}")?;
    writeln!(f, "part")?;
    writeln!(f, "{:>10}", 1)?;
    writeln!(f, "coordinates")?;

    let mut buf = [0.0f64];
    for i in 0..arr.num_tuples() {
        arr.tuple_as_f64(i, &mut buf);
        writeln!(f, "{:>12.5e}", buf[0] as f32)?;
    }

    Ok(())
}

fn write_vector_variable(path: &Path, pd: &PolyData, name: &str) -> Result<(), VtkError> {
    let arr = pd.point_data().get_array(name)
        .ok_or_else(|| VtkError::InvalidData(format!("array '{name}' not found")))?;

    let mut f = std::fs::File::create(path)?;
    writeln!(f, "{name}")?;
    writeln!(f, "part")?;
    writeln!(f, "{:>10}", 1)?;
    writeln!(f, "coordinates")?;

    let nt = arr.num_tuples();
    let mut buf = [0.0f64; 3];

    // X components
    for i in 0..nt {
        arr.tuple_as_f64(i, &mut buf);
        writeln!(f, "{:>12.5e}", buf[0] as f32)?;
    }
    // Y components
    for i in 0..nt {
        arr.tuple_as_f64(i, &mut buf);
        writeln!(f, "{:>12.5e}", buf[1] as f32)?;
    }
    // Z components
    for i in 0..nt {
        arr.tuple_as_f64(i, &mut buf);
        writeln!(f, "{:>12.5e}", buf[2] as f32)?;
    }

    Ok(())
}

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

    #[test]
    fn write_triangle() {
        let pd = PolyData::from_triangles(
            vec![[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]],
            vec![[0, 1, 2]],
        );
        let dir = std::env::temp_dir().join("vtk_ensight_test");
        let _ = std::fs::remove_dir_all(&dir);
        EnSightWriter::write(&dir, "test", &pd).unwrap();

        assert!(dir.join("test.case").exists());
        assert!(dir.join("test.geo").exists());

        let case = std::fs::read_to_string(dir.join("test.case")).unwrap();
        assert!(case.contains("ensight gold"));
        assert!(case.contains("model: test.geo"));

        let geo = std::fs::read_to_string(dir.join("test.geo")).unwrap();
        assert!(geo.contains("coordinates"));
        assert!(geo.contains("tria3"));

        let _ = std::fs::remove_dir_all(&dir);
    }

    #[test]
    fn write_with_scalars() {
        let mut pd = PolyData::from_triangles(
            vec![[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]],
            vec![[0, 1, 2]],
        );
        let s = DataArray::from_vec("temp", vec![10.0f64, 20.0, 30.0], 1);
        pd.point_data_mut().add_array(s.into());

        let dir = std::env::temp_dir().join("vtk_ensight_scalar_test");
        let _ = std::fs::remove_dir_all(&dir);
        EnSightWriter::write(&dir, "data", &pd).unwrap();

        assert!(dir.join("data.temp.scl").exists());
        let case = std::fs::read_to_string(dir.join("data.case")).unwrap();
        assert!(case.contains("scalar per node: temp"));

        let _ = std::fs::remove_dir_all(&dir);
    }
}