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

//! FreeSurfer subject container
//!
//! Provides a unified interface for loading anatomy from a FreeSurfer subject directory.

use std::collections::HashMap;
use std::path::{Path, PathBuf};

use serde::{Deserialize, Serialize};

use crate::annotation::{Annotation, annotation_filename, read_annotation};
use crate::curvature::{Curvature, curvature_filename, read_curvature};
use crate::error::{AnatomyError, Result};
use crate::surface::{Hemisphere, SurfaceMesh, SurfaceType, read_surface};

/// Handle to a loaded FreeSurfer subject
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FreeSurferHandle {
    /// Unique identifier
    pub id: String,
    /// Path to subject directory
    pub path: String,
    /// Subject ID (directory name)
    pub subject_id: String,
    /// Available surfaces
    pub available_surfaces: Vec<(Hemisphere, SurfaceType)>,
    /// Available annotations
    pub available_annotations: Vec<(Hemisphere, String)>,
}

/// A FreeSurfer subject with lazy-loaded surfaces and annotations
pub struct FreeSurferSubject {
    /// Path to subject directory
    pub path: PathBuf,
    /// Subject ID (usually the directory name)
    pub subject_id: String,
    /// Loaded surfaces (cached)
    surfaces: HashMap<(Hemisphere, SurfaceType), SurfaceMesh>,
    /// Loaded annotations (cached)
    annotations: HashMap<(Hemisphere, String), Annotation>,
    /// Loaded curvature data (cached)
    curvatures: HashMap<(Hemisphere, String), Curvature>,
}

impl FreeSurferSubject {
    /// Open a FreeSurfer subject directory
    ///
    /// # Arguments
    ///
    /// * `path` - Path to the subject directory (e.g., /subjects/fsaverage)
    ///
    /// # Returns
    ///
    /// A `FreeSurferSubject` ready for loading surfaces and annotations
    ///
    /// # Errors
    ///
    /// Returns error if the directory doesn't exist or isn't a valid FreeSurfer subject
    ///
    /// # Example
    ///
    /// ```ignore
    /// use rtx_neuro_anatomy::FreeSurferSubject;
    ///
    /// let subject = FreeSurferSubject::open("/path/to/subject")?;
    /// println!("Loaded subject: {}", subject.subject_id);
    /// ```
    pub fn open(path: impl AsRef<Path>) -> Result<Self> {
        let path = path.as_ref().to_path_buf();

        // Validate: surf directory must exist
        let surf_dir = path.join("surf");
        if !surf_dir.exists() {
            return Err(AnatomyError::InvalidSubjectDir { path });
        }

        // Extract subject ID from directory name
        let subject_id = path
            .file_name()
            .and_then(|n| n.to_str())
            .map_or_else(|| "unknown".to_string(), String::from);

        Ok(Self {
            path,
            subject_id,
            surfaces: HashMap::new(),
            annotations: HashMap::new(),
            curvatures: HashMap::new(),
        })
    }

    /// Get path to surf directory
    pub fn surf_dir(&self) -> PathBuf {
        self.path.join("surf")
    }

    /// Get path to label directory
    pub fn label_dir(&self) -> PathBuf {
        self.path.join("label")
    }

    /// Load a surface (cached)
    ///
    /// # Arguments
    ///
    /// * `hemisphere` - Left or Right hemisphere
    /// * `surface_type` - Type of surface to load
    ///
    /// # Returns
    ///
    /// Reference to the loaded surface mesh
    pub fn load_surface(
        &mut self,
        hemisphere: Hemisphere,
        surface_type: SurfaceType,
    ) -> Result<&SurfaceMesh> {
        let key = (hemisphere, surface_type);

        if !self.surfaces.contains_key(&key) {
            let filename = surface_type.filename(&hemisphere);
            let path = self.surf_dir().join(&filename);
            let mesh = read_surface(&path, surface_type, hemisphere)?;
            self.surfaces.insert(key, mesh);
        }

        Ok(self.surfaces.get(&key).unwrap())
    }

    /// Load a surface and compute normals
    pub fn load_surface_with_normals(
        &mut self,
        hemisphere: Hemisphere,
        surface_type: SurfaceType,
    ) -> Result<&SurfaceMesh> {
        self.load_surface(hemisphere, surface_type)?;

        let key = (hemisphere, surface_type);
        let mesh = self.surfaces.get_mut(&key).unwrap();
        if mesh.normals.is_none() {
            mesh.compute_normals();
        }

        Ok(self.surfaces.get(&key).unwrap())
    }

    /// Get a previously loaded surface (without loading)
    pub fn get_surface(
        &self,
        hemisphere: Hemisphere,
        surface_type: SurfaceType,
    ) -> Option<&SurfaceMesh> {
        self.surfaces.get(&(hemisphere, surface_type))
    }

    /// Get a mutable reference to a loaded surface
    pub fn get_surface_mut(
        &mut self,
        hemisphere: Hemisphere,
        surface_type: SurfaceType,
    ) -> Option<&mut SurfaceMesh> {
        self.surfaces.get_mut(&(hemisphere, surface_type))
    }

    /// Load an annotation (cached)
    ///
    /// # Arguments
    ///
    /// * `hemisphere` - Left or Right hemisphere
    /// * `atlas` - Atlas name (e.g., "aparc", "aparc.a2009s")
    ///
    /// # Returns
    ///
    /// Reference to the loaded annotation
    pub fn load_annotation(&mut self, hemisphere: Hemisphere, atlas: &str) -> Result<&Annotation> {
        let key = (hemisphere, atlas.to_string());

        if !self.annotations.contains_key(&key) {
            let filename = annotation_filename(&hemisphere, atlas);
            let path = self.label_dir().join(&filename);
            let annot = read_annotation(&path, atlas, hemisphere)?;
            self.annotations.insert(key.clone(), annot);
        }

        Ok(self.annotations.get(&key).unwrap())
    }

    /// Get a previously loaded annotation
    pub fn get_annotation(&self, hemisphere: Hemisphere, atlas: &str) -> Option<&Annotation> {
        self.annotations.get(&(hemisphere, atlas.to_string()))
    }

    /// Load curvature data (cached)
    ///
    /// # Arguments
    ///
    /// * `hemisphere` - Left or Right hemisphere
    /// * `curv_type` - Curvature type (e.g., "curv", "sulc", "thickness")
    ///
    /// # Returns
    ///
    /// Reference to the loaded curvature data
    pub fn load_curvature(
        &mut self,
        hemisphere: Hemisphere,
        curv_type: &str,
    ) -> Result<&Curvature> {
        let key = (hemisphere, curv_type.to_string());

        if !self.curvatures.contains_key(&key) {
            let filename = curvature_filename(&hemisphere, curv_type);
            let path = self.surf_dir().join(&filename);
            let curv = read_curvature(&path, curv_type, hemisphere)?;
            self.curvatures.insert(key.clone(), curv);
        }

        Ok(self.curvatures.get(&key).unwrap())
    }

    /// Get a previously loaded curvature
    pub fn get_curvature(&self, hemisphere: Hemisphere, curv_type: &str) -> Option<&Curvature> {
        self.curvatures.get(&(hemisphere, curv_type.to_string()))
    }

    /// List available surfaces in the subject directory
    pub fn list_surfaces(&self) -> Vec<(Hemisphere, SurfaceType)> {
        let mut available = Vec::new();
        let surf_dir = self.surf_dir();

        for hemi in [Hemisphere::Left, Hemisphere::Right] {
            for surf_type in [
                SurfaceType::White,
                SurfaceType::Pial,
                SurfaceType::Inflated,
                SurfaceType::Sphere,
                SurfaceType::Orig,
            ] {
                let filename = surf_type.filename(&hemi);
                if surf_dir.join(&filename).exists() {
                    available.push((hemi, surf_type));
                }
            }
        }

        available
    }

    /// List available annotations in the subject directory
    pub fn list_annotations(&self) -> Vec<(Hemisphere, String)> {
        let mut available = Vec::new();
        let label_dir = self.label_dir();

        if !label_dir.exists() {
            return available;
        }

        // Common atlases to check
        let atlases = ["aparc", "aparc.a2009s", "aparc.DKTatlas", "BA_exvivo"];

        for hemi in [Hemisphere::Left, Hemisphere::Right] {
            for atlas in &atlases {
                let filename = annotation_filename(&hemi, atlas);
                if label_dir.join(&filename).exists() {
                    available.push((hemi, atlas.to_string()));
                }
            }
        }

        available
    }

    /// List available curvature files
    pub fn list_curvatures(&self) -> Vec<(Hemisphere, String)> {
        let mut available = Vec::new();
        let surf_dir = self.surf_dir();

        let curv_types = ["curv", "sulc", "thickness", "area", "volume"];

        for hemi in [Hemisphere::Left, Hemisphere::Right] {
            for curv_type in &curv_types {
                let filename = curvature_filename(&hemi, curv_type);
                if surf_dir.join(&filename).exists() {
                    available.push((hemi, curv_type.to_string()));
                }
            }
        }

        available
    }

    /// Create a handle for IPC
    pub fn to_handle(&self, id: &str) -> FreeSurferHandle {
        FreeSurferHandle {
            id: id.to_string(),
            path: self.path.display().to_string(),
            subject_id: self.subject_id.clone(),
            available_surfaces: self.list_surfaces(),
            available_annotations: self.list_annotations(),
        }
    }

    /// Unload all cached data to free memory
    pub fn clear_cache(&mut self) {
        self.surfaces.clear();
        self.annotations.clear();
        self.curvatures.clear();
    }

    /// Number of loaded surfaces
    pub fn n_loaded_surfaces(&self) -> usize {
        self.surfaces.len()
    }

    /// Number of loaded annotations
    pub fn n_loaded_annotations(&self) -> usize {
        self.annotations.len()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::fs;
    use tempfile::TempDir;

    fn create_test_subject() -> (TempDir, PathBuf) {
        let dir = TempDir::new().unwrap();
        let subject_path = dir.path().join("test_subject");

        // Create required directories
        fs::create_dir_all(subject_path.join("surf")).unwrap();
        fs::create_dir_all(subject_path.join("label")).unwrap();

        (dir, subject_path)
    }

    #[test]
    fn test_open_subject() {
        let (_dir, subject_path) = create_test_subject();
        let subject = FreeSurferSubject::open(&subject_path).unwrap();

        assert_eq!(subject.subject_id, "test_subject");
        assert_eq!(subject.path, subject_path);
    }

    #[test]
    fn test_open_invalid_directory() {
        let result = FreeSurferSubject::open("/nonexistent/path");
        assert!(matches!(
            result,
            Err(AnatomyError::InvalidSubjectDir { .. })
        ));
    }

    #[test]
    fn test_list_empty_subject() {
        let (_dir, subject_path) = create_test_subject();
        let subject = FreeSurferSubject::open(&subject_path).unwrap();

        // No surfaces or annotations should be found
        assert!(subject.list_surfaces().is_empty());
        assert!(subject.list_annotations().is_empty());
    }

    #[test]
    fn test_handle_creation() {
        let (_dir, subject_path) = create_test_subject();
        let subject = FreeSurferSubject::open(&subject_path).unwrap();

        let handle = subject.to_handle("test-id");
        assert_eq!(handle.id, "test-id");
        assert_eq!(handle.subject_id, "test_subject");
    }

    #[test]
    fn test_clear_cache() {
        let (_dir, subject_path) = create_test_subject();
        let mut subject = FreeSurferSubject::open(&subject_path).unwrap();

        assert_eq!(subject.n_loaded_surfaces(), 0);
        subject.clear_cache();
        assert_eq!(subject.n_loaded_surfaces(), 0);
    }
}