rustytorch/crates/specialized/rtx-neuro-forward/src/source_space.rs

//! Source space definitions for source localization.
//!
//! A source space defines the possible locations and orientations
//! of current dipoles in the brain.

use crate::{Orientation, Position, norm, normalize};
use nalgebra::Vector3;
use std::f64::consts::PI;

/// Orientation constraint for sources
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum SourceOrientation {
    /// Free orientation (3 DOF per source)
    Free,
    /// Fixed orientation perpendicular to cortical surface
    Fixed,
    /// Loose constraint (partially constrained)
    Loose(f64),
}

/// A single source point (dipole location)
#[derive(Debug, Clone)]
pub struct SourcePoint {
    /// Position in meters
    position: Position,
    /// Orientation (if fixed)
    orientation: Option<Orientation>,
    /// Surface normal (for cortical sources)
    normal: Option<Orientation>,
    /// Hemisphere (left=-1, right=1)
    hemisphere: i8,
    /// Vertex index (for surface sources)
    vertex_index: Option<usize>,
}

impl SourcePoint {
    /// Create a new source point with free orientation
    pub fn new(position: [f64; 3]) -> Self {
        Self {
            position: Vector3::new(position[0], position[1], position[2]),
            orientation: None,
            normal: None,
            hemisphere: 0,
            vertex_index: None,
        }
    }

    /// Create a source point with fixed orientation
    pub fn with_orientation(position: [f64; 3], orientation: [f64; 3]) -> Self {
        let ori = Vector3::new(orientation[0], orientation[1], orientation[2]);
        Self {
            position: Vector3::new(position[0], position[1], position[2]),
            orientation: Some(normalize(&ori)),
            normal: Some(normalize(&ori)),
            hemisphere: 0,
            vertex_index: None,
        }
    }

    /// Create a cortical source with surface normal
    pub fn cortical(
        position: [f64; 3],
        normal: [f64; 3],
        hemisphere: i8,
        vertex_index: usize,
    ) -> Self {
        let n = Vector3::new(normal[0], normal[1], normal[2]);
        Self {
            position: Vector3::new(position[0], position[1], position[2]),
            orientation: Some(normalize(&n)),
            normal: Some(normalize(&n)),
            hemisphere,
            vertex_index: Some(vertex_index),
        }
    }

    /// Get the source position
    pub fn position(&self) -> &Position {
        &self.position
    }

    /// Get the source orientation (if fixed)
    pub fn orientation(&self) -> Option<Orientation> {
        self.orientation
    }

    /// Get the surface normal
    pub fn normal(&self) -> Option<&Orientation> {
        self.normal.as_ref()
    }

    /// Get hemisphere (-1=left, 1=right, 0=unknown)
    pub fn hemisphere(&self) -> i8 {
        self.hemisphere
    }

    /// Get vertex index for surface sources
    pub fn vertex_index(&self) -> Option<usize> {
        self.vertex_index
    }
}

/// Source space containing all source locations
#[derive(Debug, Clone)]
pub struct SourceSpace {
    /// Source points
    sources: Vec<SourcePoint>,
    /// Orientation constraint
    orientation: SourceOrientation,
    /// Subject name
    subject: Option<String>,
}

impl SourceSpace {
    /// Create a new empty source space
    pub fn new(orientation: SourceOrientation) -> Self {
        Self {
            sources: Vec::new(),
            orientation,
            subject: None,
        }
    }

    /// Create a source space from points
    pub fn from_points(points: Vec<SourcePoint>, orientation: SourceOrientation) -> Self {
        Self {
            sources: points,
            orientation,
            subject: None,
        }
    }

    /// Add a source point
    pub fn add_source(&mut self, source: SourcePoint) {
        self.sources.push(source);
    }

    /// Get number of sources
    pub fn len(&self) -> usize {
        self.sources.len()
    }

    /// Check if empty
    pub fn is_empty(&self) -> bool {
        self.sources.is_empty()
    }

    /// Check if sources have fixed orientation
    pub fn is_fixed_orientation(&self) -> bool {
        matches!(self.orientation, SourceOrientation::Fixed)
    }

    /// Get orientation constraint
    pub fn orientation_constraint(&self) -> SourceOrientation {
        self.orientation
    }

    /// Iterate over sources
    pub fn iter(&self) -> impl Iterator<Item = &SourcePoint> {
        self.sources.iter()
    }

    /// Get a specific source
    pub fn get(&self, index: usize) -> Option<&SourcePoint> {
        self.sources.get(index)
    }

    /// Create a volume source space (regular grid)
    ///
    /// # Arguments
    /// * `bounds` - Bounding box [(xmin, xmax), (ymin, ymax), (zmin, zmax)]
    /// * `spacing` - Grid spacing in meters
    /// * `exclude_radius` - Optional radius from origin to exclude (for ventricles)
    pub fn create_volume_grid(
        bounds: [(f64, f64); 3],
        spacing: f64,
        exclude_radius: Option<f64>,
    ) -> Self {
        let mut sources = Vec::new();

        let mut x = bounds[0].0;
        while x <= bounds[0].1 {
            let mut y = bounds[1].0;
            while y <= bounds[1].1 {
                let mut z = bounds[2].0;
                while z <= bounds[2].1 {
                    let pos = Vector3::new(x, y, z);

                    // Check exclusion radius
                    if let Some(r) = exclude_radius
                        && norm(&pos) < r
                    {
                        z += spacing;
                        continue;
                    }

                    sources.push(SourcePoint::new([x, y, z]));
                    z += spacing;
                }
                y += spacing;
            }
            x += spacing;
        }

        Self {
            sources,
            orientation: SourceOrientation::Free,
            subject: None,
        }
    }

    /// Create a spherical shell source space
    ///
    /// # Arguments
    /// * `center` - Center of the sphere
    /// * `radius` - Radius of the shell
    /// * `n_points` - Approximate number of points
    pub fn create_spherical_shell(center: [f64; 3], radius: f64, n_points: usize) -> Self {
        let mut sources = Vec::new();
        let center = Vector3::new(center[0], center[1], center[2]);

        // Use Fibonacci sphere for uniform distribution
        let golden_ratio = f64::midpoint(1.0, 5.0_f64.sqrt());

        for i in 0..n_points {
            let theta = 2.0 * PI * i as f64 / golden_ratio;
            let phi = (1.0 - 2.0 * (i as f64 + 0.5) / n_points as f64).acos();

            let x = center.x + radius * phi.sin() * theta.cos();
            let y = center.y + radius * phi.sin() * theta.sin();
            let z = center.z + radius * phi.cos();

            // Normal points outward
            let normal = normalize(&Vector3::new(x - center.x, y - center.y, z - center.z));

            let hemisphere = if x < center.x { -1 } else { 1 };

            sources.push(SourcePoint {
                position: Vector3::new(x, y, z),
                orientation: Some(normal),
                normal: Some(normal),
                hemisphere,
                vertex_index: Some(i),
            });
        }

        Self {
            sources,
            orientation: SourceOrientation::Fixed,
            subject: None,
        }
    }

    /// Create a source space on two hemispheric shells (simple cortex model)
    pub fn create_cortex_shells(center: [f64; 3], radius: f64, n_per_hemisphere: usize) -> Self {
        let mut sources = Vec::new();
        let center = Vector3::new(center[0], center[1], center[2]);
        let golden_ratio = f64::midpoint(1.0, 5.0_f64.sqrt());

        // Left hemisphere
        for i in 0..n_per_hemisphere {
            let theta = 2.0 * PI * i as f64 / golden_ratio;
            let phi = (1.0 - 2.0 * (i as f64 + 0.5) / (2 * n_per_hemisphere) as f64).acos();

            let x = center.x - radius * phi.sin() * theta.cos().abs();
            let y = center.y + radius * phi.sin() * theta.sin();
            let z = center.z + radius * phi.cos();

            let normal = normalize(&Vector3::new(x - center.x, y - center.y, z - center.z));

            sources.push(SourcePoint {
                position: Vector3::new(x, y, z),
                orientation: Some(normal),
                normal: Some(normal),
                hemisphere: -1,
                vertex_index: Some(i),
            });
        }

        // Right hemisphere
        for i in 0..n_per_hemisphere {
            let theta = 2.0 * PI * i as f64 / golden_ratio;
            let phi = (1.0 - 2.0 * (i as f64 + 0.5) / (2 * n_per_hemisphere) as f64).acos();

            let x = center.x + radius * phi.sin() * theta.cos().abs();
            let y = center.y + radius * phi.sin() * theta.sin();
            let z = center.z + radius * phi.cos();

            let normal = normalize(&Vector3::new(x - center.x, y - center.y, z - center.z));

            sources.push(SourcePoint {
                position: Vector3::new(x, y, z),
                orientation: Some(normal),
                normal: Some(normal),
                hemisphere: 1,
                vertex_index: Some(n_per_hemisphere + i),
            });
        }

        Self {
            sources,
            orientation: SourceOrientation::Fixed,
            subject: None,
        }
    }
}

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

    #[test]
    fn test_source_point() {
        let src = SourcePoint::new([0.01, 0.02, 0.03]);
        assert!((src.position().x - 0.01).abs() < 1e-10);
    }

    #[test]
    fn test_source_with_orientation() {
        let src = SourcePoint::with_orientation([0.0, 0.0, 0.05], [0.0, 0.0, 1.0]);
        let ori = src.orientation().unwrap();
        assert!((ori.z - 1.0).abs() < 1e-10);
    }

    #[test]
    fn test_volume_grid() {
        let ss = SourceSpace::create_volume_grid(
            [(-0.02, 0.02), (-0.02, 0.02), (-0.02, 0.02)],
            0.01,
            None,
        );
        assert!(ss.len() > 0);
    }

    #[test]
    fn test_spherical_shell() {
        let ss = SourceSpace::create_spherical_shell([0.0, 0.0, 0.04], 0.06, 100);
        assert_eq!(ss.len(), 100);
        assert!(ss.is_fixed_orientation());
    }

    #[test]
    fn test_cortex_shells() {
        let ss = SourceSpace::create_cortex_shells([0.0, 0.0, 0.04], 0.06, 50);
        assert_eq!(ss.len(), 100); // 50 per hemisphere
    }
}