rustytorch/crates/specialized/rtx-neuro-anatomy/src/surface.rs

//! FreeSurfer surface file reading
//!
//! Reads binary surface files (.pial, .white, .inflated, .sphere, .orig)
//! in the FreeSurfer triangle surface format.

use std::fs::File;
use std::io::{BufReader, Read};
use std::path::Path;

use byteorder::{BigEndian, ReadBytesExt};
use nalgebra::Vector3;
use serde::{Deserialize, Serialize};

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

/// FreeSurfer triangle surface magic number
const TRIANGLE_SURFACE_MAGIC: [u8; 3] = [0xFF, 0xFF, 0xFE];

/// Type of brain surface
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum SurfaceType {
    /// White matter surface (inner boundary)
    White,
    /// Pial surface (outer boundary, follows gyri/sulci)
    Pial,
    /// Inflated surface (smoothed for visualization)
    Inflated,
    /// Spherical surface (for registration)
    Sphere,
    /// Original surface (from segmentation)
    Orig,
}

impl SurfaceType {
    /// Get the filename for this surface type
    pub fn filename(&self, hemisphere: &Hemisphere) -> String {
        let prefix = hemisphere.prefix();
        let suffix = match self {
            SurfaceType::White => "white",
            SurfaceType::Pial => "pial",
            SurfaceType::Inflated => "inflated",
            SurfaceType::Sphere => "sphere",
            SurfaceType::Orig => "orig",
        };
        format!("{}.{}", prefix, suffix)
    }
}

/// Brain hemisphere
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum Hemisphere {
    /// Left hemisphere
    Left,
    /// Right hemisphere
    Right,
}

impl Hemisphere {
    /// Get the filename prefix for this hemisphere
    pub fn prefix(&self) -> &str {
        match self {
            Hemisphere::Left => "lh",
            Hemisphere::Right => "rh",
        }
    }
}

/// A triangulated surface mesh
#[derive(Debug, Clone)]
pub struct SurfaceMesh {
    /// Vertex positions [n_vertices] (x, y, z in mm, RAS coordinates)
    pub vertices: Vec<Vector3<f32>>,
    /// Triangle face indices [n_faces] (3 vertex indices per face)
    pub faces: Vec<[u32; 3]>,
    /// Surface type
    pub surface_type: SurfaceType,
    /// Hemisphere
    pub hemisphere: Hemisphere,
    /// Vertex normals (computed on demand)
    pub normals: Option<Vec<Vector3<f32>>>,
}

impl SurfaceMesh {
    /// Number of vertices
    pub fn n_vertices(&self) -> usize {
        self.vertices.len()
    }

    /// Number of faces (triangles)
    pub fn n_faces(&self) -> usize {
        self.faces.len()
    }

    /// Compute vertex normals from face normals
    ///
    /// Uses area-weighted averaging of adjacent face normals
    pub fn compute_normals(&mut self) {
        let mut normals: Vec<Vector3<f32>> = vec![Vector3::zeros(); self.vertices.len()];

        // Accumulate face normals at each vertex
        for face in &self.faces {
            let v0 = self.vertices[face[0] as usize];
            let v1 = self.vertices[face[1] as usize];
            let v2 = self.vertices[face[2] as usize];

            // Compute face normal (cross product)
            let edge1 = v1 - v0;
            let edge2 = v2 - v0;
            let face_normal = edge1.cross(&edge2);

            // Accumulate (area-weighted by not normalizing the cross product)
            normals[face[0] as usize] += face_normal;
            normals[face[1] as usize] += face_normal;
            normals[face[2] as usize] += face_normal;
        }

        // Normalize all vertex normals
        for normal in &mut normals {
            let len = normal.norm();
            if len > 1e-10 {
                *normal /= len;
            }
        }

        self.normals = Some(normals);
    }

    /// Get vertices as flat array [[x, y, z], ...]
    pub fn vertices_flat(&self) -> Vec<[f32; 3]> {
        self.vertices.iter().map(|v| [v.x, v.y, v.z]).collect()
    }

    /// Get normals as flat array (computes if not available)
    pub fn normals_flat(&mut self) -> Vec<[f32; 3]> {
        if self.normals.is_none() {
            self.compute_normals();
        }
        self.normals
            .as_ref()
            .unwrap()
            .iter()
            .map(|n| [n.x, n.y, n.z])
            .collect()
    }

    /// Get bounding box [min, max]
    pub fn bounds(&self) -> ([f32; 3], [f32; 3]) {
        let mut min = [f32::MAX; 3];
        let mut max = [f32::MIN; 3];

        for v in &self.vertices {
            min[0] = min[0].min(v.x);
            min[1] = min[1].min(v.y);
            min[2] = min[2].min(v.z);
            max[0] = max[0].max(v.x);
            max[1] = max[1].max(v.y);
            max[2] = max[2].max(v.z);
        }

        (min, max)
    }

    /// Get center of mass
    pub fn center(&self) -> Vector3<f32> {
        let mut sum = Vector3::zeros();
        for v in &self.vertices {
            sum += v;
        }
        sum / self.vertices.len() as f32
    }
}

/// Read a FreeSurfer surface file
///
/// # Arguments
///
/// * `path` - Path to the surface file
/// * `surface_type` - Type of surface being read
/// * `hemisphere` - Hemisphere of the surface
///
/// # Returns
///
/// A `SurfaceMesh` containing the vertices and faces
///
/// # Example
///
/// ```ignore
/// use rtx_neuro_anatomy::surface::{read_surface, SurfaceType, Hemisphere};
///
/// let mesh = read_surface(
///     "/path/to/subject/surf/lh.pial",
///     SurfaceType::Pial,
///     Hemisphere::Left,
/// )?;
/// println!("Loaded {} vertices, {} faces", mesh.n_vertices(), mesh.n_faces());
/// ```
pub fn read_surface(
    path: impl AsRef<Path>,
    surface_type: SurfaceType,
    hemisphere: Hemisphere,
) -> Result<SurfaceMesh> {
    let path = path.as_ref();
    let file = File::open(path).map_err(|_| AnatomyError::SurfaceNotFound {
        path: path.to_path_buf(),
    })?;
    let mut reader = BufReader::new(file);

    // Read magic number (3 bytes, big-endian)
    let mut magic = [0u8; 3];
    reader.read_exact(&mut magic)?;

    if magic != TRIANGLE_SURFACE_MAGIC {
        return Err(AnatomyError::InvalidMagic {
            expected: TRIANGLE_SURFACE_MAGIC.to_vec(),
            found: magic.to_vec(),
        });
    }

    // Skip comment (read until two consecutive newlines)
    skip_comment(&mut reader)?;

    // Read vertex and face counts
    let n_vertices = reader.read_i32::<BigEndian>()? as usize;
    let n_faces = reader.read_i32::<BigEndian>()? as usize;

    // Read vertices (n_vertices * 3 floats, big-endian)
    let mut vertices = Vec::with_capacity(n_vertices);
    for _ in 0..n_vertices {
        let x = reader.read_f32::<BigEndian>()?;
        let y = reader.read_f32::<BigEndian>()?;
        let z = reader.read_f32::<BigEndian>()?;
        vertices.push(Vector3::new(x, y, z));
    }

    // Read faces (n_faces * 3 ints, big-endian)
    let mut faces = Vec::with_capacity(n_faces);
    for _ in 0..n_faces {
        let v0 = reader.read_i32::<BigEndian>()? as u32;
        let v1 = reader.read_i32::<BigEndian>()? as u32;
        let v2 = reader.read_i32::<BigEndian>()? as u32;
        faces.push([v0, v1, v2]);
    }

    Ok(SurfaceMesh {
        vertices,
        faces,
        surface_type,
        hemisphere,
        normals: None,
    })
}

/// Skip the comment section in a FreeSurfer surface file
///
/// The comment ends at two consecutive newlines
fn skip_comment(reader: &mut BufReader<File>) -> Result<()> {
    let mut prev = 0u8;
    let mut buf = [0u8; 1];

    loop {
        reader.read_exact(&mut buf)?;
        if prev == b'\n' && buf[0] == b'\n' {
            break;
        }
        prev = buf[0];
    }

    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::io::Write;
    use tempfile::NamedTempFile;

    /// Create a minimal test surface file
    fn create_test_surface() -> NamedTempFile {
        use byteorder::WriteBytesExt;

        let mut file = NamedTempFile::new().unwrap();

        // Write magic
        file.write_all(&TRIANGLE_SURFACE_MAGIC).unwrap();

        // Write comment (with double newline)
        file.write_all(b"test surface\n\n").unwrap();

        // Write counts (4 vertices, 2 faces forming a pyramid base)
        file.write_i32::<BigEndian>(4).unwrap();
        file.write_i32::<BigEndian>(2).unwrap();

        // Write vertices (a square)
        // v0: (0, 0, 0)
        file.write_f32::<BigEndian>(0.0).unwrap();
        file.write_f32::<BigEndian>(0.0).unwrap();
        file.write_f32::<BigEndian>(0.0).unwrap();
        // v1: (1, 0, 0)
        file.write_f32::<BigEndian>(1.0).unwrap();
        file.write_f32::<BigEndian>(0.0).unwrap();
        file.write_f32::<BigEndian>(0.0).unwrap();
        // v2: (1, 1, 0)
        file.write_f32::<BigEndian>(1.0).unwrap();
        file.write_f32::<BigEndian>(1.0).unwrap();
        file.write_f32::<BigEndian>(0.0).unwrap();
        // v3: (0, 1, 0)
        file.write_f32::<BigEndian>(0.0).unwrap();
        file.write_f32::<BigEndian>(1.0).unwrap();
        file.write_f32::<BigEndian>(0.0).unwrap();

        // Write faces (two triangles)
        // Face 0: (0, 1, 2)
        file.write_i32::<BigEndian>(0).unwrap();
        file.write_i32::<BigEndian>(1).unwrap();
        file.write_i32::<BigEndian>(2).unwrap();
        // Face 1: (0, 2, 3)
        file.write_i32::<BigEndian>(0).unwrap();
        file.write_i32::<BigEndian>(2).unwrap();
        file.write_i32::<BigEndian>(3).unwrap();

        file.flush().unwrap();
        file
    }

    #[test]
    fn test_read_surface() {
        let file = create_test_surface();
        let mesh = read_surface(file.path(), SurfaceType::Pial, Hemisphere::Left).unwrap();

        assert_eq!(mesh.n_vertices(), 4);
        assert_eq!(mesh.n_faces(), 2);
        assert_eq!(mesh.surface_type, SurfaceType::Pial);
        assert_eq!(mesh.hemisphere, Hemisphere::Left);

        // Check first vertex
        assert!((mesh.vertices[0].x - 0.0).abs() < 1e-6);
        assert!((mesh.vertices[0].y - 0.0).abs() < 1e-6);
        assert!((mesh.vertices[0].z - 0.0).abs() < 1e-6);

        // Check second vertex
        assert!((mesh.vertices[1].x - 1.0).abs() < 1e-6);
    }

    #[test]
    fn test_compute_normals() {
        let file = create_test_surface();
        let mut mesh = read_surface(file.path(), SurfaceType::Pial, Hemisphere::Left).unwrap();

        mesh.compute_normals();
        assert!(mesh.normals.is_some());

        let normals = mesh.normals.as_ref().unwrap();
        assert_eq!(normals.len(), 4);

        // All normals should point in +Z direction for a flat XY plane
        for normal in normals {
            assert!((normal.z - 1.0).abs() < 1e-6);
        }
    }

    #[test]
    fn test_bounds() {
        let file = create_test_surface();
        let mesh = read_surface(file.path(), SurfaceType::Pial, Hemisphere::Left).unwrap();

        let (min, max) = mesh.bounds();
        assert!((min[0] - 0.0).abs() < 1e-6);
        assert!((min[1] - 0.0).abs() < 1e-6);
        assert!((min[2] - 0.0).abs() < 1e-6);
        assert!((max[0] - 1.0).abs() < 1e-6);
        assert!((max[1] - 1.0).abs() < 1e-6);
        assert!((max[2] - 0.0).abs() < 1e-6);
    }

    #[test]
    fn test_center() {
        let file = create_test_surface();
        let mesh = read_surface(file.path(), SurfaceType::Pial, Hemisphere::Left).unwrap();

        let center = mesh.center();
        assert!((center.x - 0.5).abs() < 1e-6);
        assert!((center.y - 0.5).abs() < 1e-6);
        assert!((center.z - 0.0).abs() < 1e-6);
    }

    #[test]
    fn test_surface_type_filename() {
        assert_eq!(SurfaceType::Pial.filename(&Hemisphere::Left), "lh.pial");
        assert_eq!(SurfaceType::White.filename(&Hemisphere::Right), "rh.white");
        assert_eq!(
            SurfaceType::Inflated.filename(&Hemisphere::Left),
            "lh.inflated"
        );
    }

    #[test]
    fn test_invalid_magic() {
        use std::io::Write;
        let mut file = NamedTempFile::new().unwrap();
        file.write_all(&[0x00, 0x00, 0x00]).unwrap();
        file.flush().unwrap();

        let result = read_surface(file.path(), SurfaceType::Pial, Hemisphere::Left);
        assert!(matches!(result, Err(AnatomyError::InvalidMagic { .. })));
    }
}