rustytorch/crates/specialized/rtx-neuro-pinn/src/conductivity.rs

//! Learnable tissue conductivity models for head modeling.
//!
//! Provides models for tissue conductivity that can be learned
//! jointly with source localization or used with fixed values.

use crate::error::{PinnError, PinnResult};
use ndarray::{Array1, Array2};
use serde::{Deserialize, Serialize};

/// Standard tissue types in the head
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum TissueType {
    /// Brain gray matter
    GrayMatter,
    /// Brain white matter
    WhiteMatter,
    /// Cerebrospinal fluid
    Csf,
    /// Skull bone
    Skull,
    /// Scalp/skin
    Scalp,
    /// Air (e.g., sinuses)
    Air,
    /// Eye tissue
    Eye,
    /// Muscle tissue
    Muscle,
}

impl TissueType {
    /// Default isotropic conductivity (S/m)
    pub fn default_conductivity(&self) -> f64 {
        match self {
            TissueType::GrayMatter => 0.33,
            TissueType::WhiteMatter => 0.14,
            TissueType::Csf => 1.79,
            TissueType::Skull => 0.0042,
            TissueType::Scalp => 0.43,
            TissueType::Air => 1e-12,
            TissueType::Eye => 1.5,
            TissueType::Muscle => 0.2,
        }
    }

    /// Conductivity range (min, max) for learning
    pub fn conductivity_range(&self) -> (f64, f64) {
        match self {
            TissueType::GrayMatter => (0.1, 0.6),
            TissueType::WhiteMatter => (0.05, 0.3),
            TissueType::Csf => (1.0, 2.5),
            TissueType::Skull => (0.001, 0.02),
            TissueType::Scalp => (0.2, 0.7),
            TissueType::Air => (1e-15, 1e-10),
            TissueType::Eye => (0.5, 2.5),
            TissueType::Muscle => (0.1, 0.4),
        }
    }
}

/// A tissue layer with conductivity
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TissueLayer {
    /// Tissue type
    pub tissue_type: TissueType,
    /// Name/label
    pub name: String,
    /// Isotropic conductivity (S/m)
    pub conductivity: f64,
    /// Inner radius (for spherical models)
    pub inner_radius: Option<f64>,
    /// Outer radius (for spherical models)
    pub outer_radius: Option<f64>,
    /// Whether conductivity is learnable
    pub learnable: bool,
}

impl TissueLayer {
    /// Create a new tissue layer
    pub fn new(tissue_type: TissueType, name: impl Into<String>) -> Self {
        Self {
            tissue_type,
            name: name.into(),
            conductivity: tissue_type.default_conductivity(),
            inner_radius: None,
            outer_radius: None,
            learnable: false,
        }
    }

    /// Set conductivity
    pub fn with_conductivity(mut self, sigma: f64) -> Self {
        self.conductivity = sigma;
        self
    }

    /// Set radii for spherical model
    pub fn with_radii(mut self, inner: f64, outer: f64) -> Self {
        self.inner_radius = Some(inner);
        self.outer_radius = Some(outer);
        self
    }

    /// Make learnable
    pub fn learnable(mut self) -> Self {
        self.learnable = true;
        self
    }

    /// Check if point is inside this layer (spherical model)
    pub fn contains(&self, r: f64) -> bool {
        match (self.inner_radius, self.outer_radius) {
            (Some(inner), Some(outer)) => r >= inner && r < outer,
            (Some(inner), None) => r >= inner,
            (None, Some(outer)) => r < outer,
            (None, None) => true,
        }
    }
}

/// Conductivity model trait
pub trait ConductivityModel {
    /// Get conductivity at a point
    fn conductivity_at(&self, point: &[f64; 3]) -> f64;

    /// Get conductivity field at multiple points
    fn conductivity_field(&self, points: &Array2<f64>) -> Array1<f64> {
        let n = points.nrows();
        let mut field = Array1::zeros(n);
        for i in 0..n {
            field[i] = self.conductivity_at(&[points[[i, 0]], points[[i, 1]], points[[i, 2]]]);
        }
        field
    }

    /// Whether this model has learnable parameters
    fn is_learnable(&self) -> bool {
        false
    }

    /// Get learnable parameters
    fn get_parameters(&self) -> Vec<f64> {
        Vec::new()
    }

    /// Set learnable parameters
    fn set_parameters(&mut self, _params: &[f64]) -> PinnResult<()> {
        Ok(())
    }
}

/// Homogeneous (constant) conductivity
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HomogeneousConductivity {
    /// Conductivity value (S/m)
    sigma: f64,
    /// Whether learnable
    learnable: bool,
}

impl HomogeneousConductivity {
    /// Create with fixed conductivity
    pub fn new(sigma: f64) -> Self {
        Self {
            sigma,
            learnable: false,
        }
    }

    /// Create with learnable conductivity
    pub fn learnable(initial: f64) -> Self {
        Self {
            sigma: initial,
            learnable: true,
        }
    }

    /// Get conductivity value
    pub fn sigma(&self) -> f64 {
        self.sigma
    }
}

impl ConductivityModel for HomogeneousConductivity {
    fn conductivity_at(&self, _point: &[f64; 3]) -> f64 {
        self.sigma
    }

    fn is_learnable(&self) -> bool {
        self.learnable
    }

    fn get_parameters(&self) -> Vec<f64> {
        if self.learnable {
            vec![self.sigma]
        } else {
            Vec::new()
        }
    }

    fn set_parameters(&mut self, params: &[f64]) -> PinnResult<()> {
        if self.learnable {
            if params.is_empty() {
                return Err(PinnError::DimensionMismatch(
                    "Expected 1 parameter for learnable conductivity".into(),
                ));
            }
            self.sigma = params[0].max(1e-6); // Ensure positive
        }
        Ok(())
    }
}

/// Layered spherical conductivity model
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LayeredConductivity {
    /// Tissue layers (ordered from inside to outside)
    layers: Vec<TissueLayer>,
    /// Center of the spherical model
    center: [f64; 3],
}

impl LayeredConductivity {
    /// Create a new layered model
    pub fn new(center: [f64; 3]) -> Self {
        Self {
            layers: Vec::new(),
            center,
        }
    }

    /// Add a layer
    pub fn add_layer(&mut self, layer: TissueLayer) {
        self.layers.push(layer);
        // Sort by outer radius
        self.layers.sort_by(|a, b| {
            let ra = a.outer_radius.unwrap_or(f64::MAX);
            let rb = b.outer_radius.unwrap_or(f64::MAX);
            ra.partial_cmp(&rb).unwrap()
        });
    }

    /// Create standard 3-layer head model
    pub fn three_layer(
        center: [f64; 3],
        brain_radius: f64,
        skull_thickness: f64,
        scalp_thickness: f64,
    ) -> Self {
        let skull_outer = brain_radius + skull_thickness;
        let scalp_outer = skull_outer + scalp_thickness;

        let mut model = Self::new(center);

        model.add_layer(
            TissueLayer::new(TissueType::GrayMatter, "brain").with_radii(0.0, brain_radius),
        );

        model.add_layer(
            TissueLayer::new(TissueType::Skull, "skull").with_radii(brain_radius, skull_outer),
        );

        model.add_layer(
            TissueLayer::new(TissueType::Scalp, "scalp").with_radii(skull_outer, scalp_outer),
        );

        model
    }

    /// Create standard 4-layer head model (with CSF)
    pub fn four_layer(
        center: [f64; 3],
        brain_radius: f64,
        csf_thickness: f64,
        skull_thickness: f64,
        scalp_thickness: f64,
    ) -> Self {
        let csf_outer = brain_radius + csf_thickness;
        let skull_outer = csf_outer + skull_thickness;
        let scalp_outer = skull_outer + scalp_thickness;

        let mut model = Self::new(center);

        model.add_layer(
            TissueLayer::new(TissueType::GrayMatter, "brain").with_radii(0.0, brain_radius),
        );

        model.add_layer(
            TissueLayer::new(TissueType::Csf, "csf").with_radii(brain_radius, csf_outer),
        );

        model.add_layer(
            TissueLayer::new(TissueType::Skull, "skull").with_radii(csf_outer, skull_outer),
        );

        model.add_layer(
            TissueLayer::new(TissueType::Scalp, "scalp").with_radii(skull_outer, scalp_outer),
        );

        model
    }

    /// Get layer containing a point
    pub fn layer_at(&self, point: &[f64; 3]) -> Option<&TissueLayer> {
        let dx = point[0] - self.center[0];
        let dy = point[1] - self.center[1];
        let dz = point[2] - self.center[2];
        let r = (dx * dx + dy * dy + dz * dz).sqrt();

        self.layers
            .iter()
            .find(|&layer| layer.contains(r))
            .map(|v| v as _)
    }

    /// Get all layers
    pub fn layers(&self) -> &[TissueLayer] {
        &self.layers
    }

    /// Get mutable layers
    pub fn layers_mut(&mut self) -> &mut [TissueLayer] {
        &mut self.layers
    }

    /// Number of layers
    pub fn n_layers(&self) -> usize {
        self.layers.len()
    }
}

impl ConductivityModel for LayeredConductivity {
    fn conductivity_at(&self, point: &[f64; 3]) -> f64 {
        self.layer_at(point).map_or(0.0, |l| l.conductivity)
    }

    fn is_learnable(&self) -> bool {
        self.layers.iter().any(|l| l.learnable)
    }

    fn get_parameters(&self) -> Vec<f64> {
        self.layers
            .iter()
            .filter(|l| l.learnable)
            .map(|l| l.conductivity)
            .collect()
    }

    fn set_parameters(&mut self, params: &[f64]) -> PinnResult<()> {
        let learnable_count = self.layers.iter().filter(|l| l.learnable).count();
        if params.len() != learnable_count {
            return Err(PinnError::DimensionMismatch(format!(
                "Expected {} parameters, got {}",
                learnable_count,
                params.len()
            )));
        }

        let mut idx = 0;
        for layer in &mut self.layers {
            if layer.learnable {
                let (min, max) = layer.tissue_type.conductivity_range();
                layer.conductivity = params[idx].clamp(min, max);
                idx += 1;
            }
        }

        Ok(())
    }
}

/// Neural network-based learnable conductivity
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LearnableConductivity {
    /// Base layered model
    base: LayeredConductivity,
    /// Perturbation weights [n_basis, n_layers]
    perturbation_weights: Array2<f64>,
    /// Radial basis function centers
    rbf_centers: Array1<f64>,
    /// RBF width
    rbf_width: f64,
    /// Maximum perturbation magnitude
    max_perturbation: f64,
}

impl LearnableConductivity {
    /// Create learnable conductivity from base model
    pub fn new(base: LayeredConductivity, n_basis: usize) -> Self {
        let n_layers = base.n_layers();

        // Distribute RBF centers radially
        let max_radius = base
            .layers()
            .iter()
            .filter_map(|l| l.outer_radius)
            .fold(0.0f64, f64::max);

        let mut rbf_centers = Array1::zeros(n_basis);
        for i in 0..n_basis {
            rbf_centers[i] = max_radius * (i as f64 + 0.5) / n_basis as f64;
        }

        Self {
            base,
            perturbation_weights: Array2::zeros((n_basis, n_layers)),
            rbf_centers,
            rbf_width: max_radius / (n_basis as f64 * 2.0),
            max_perturbation: 0.1, // 10% max perturbation
        }
    }

    /// Compute RBF values at radius
    fn rbf_values(&self, r: f64) -> Array1<f64> {
        self.rbf_centers.mapv(|c| {
            let d = (r - c) / self.rbf_width;
            (-0.5 * d * d).exp()
        })
    }

    /// Get perturbation at radius for each layer
    fn perturbation(&self, r: f64) -> Array1<f64> {
        let rbf = self.rbf_values(r);
        let mut perturb: Array1<f64> = Array1::zeros(self.base.n_layers());

        for j in 0..self.base.n_layers() {
            for i in 0..self.rbf_centers.len() {
                perturb[j] += self.perturbation_weights[[i, j]] * rbf[i];
            }
            // Clamp perturbation
            perturb[j] = perturb[j].tanh() * self.max_perturbation;
        }

        perturb
    }

    /// Get base model
    pub fn base(&self) -> &LayeredConductivity {
        &self.base
    }
}

impl ConductivityModel for LearnableConductivity {
    fn conductivity_at(&self, point: &[f64; 3]) -> f64 {
        let dx = point[0] - self.base.center[0];
        let dy = point[1] - self.base.center[1];
        let dz = point[2] - self.base.center[2];
        let r = (dx * dx + dy * dy + dz * dz).sqrt();

        // Find which layer and apply perturbation
        let perturb = self.perturbation(r);

        for (i, layer) in self.base.layers().iter().enumerate() {
            if layer.contains(r) {
                let (min, max) = layer.tissue_type.conductivity_range();
                let sigma = layer.conductivity * (1.0 + perturb[i]);
                return sigma.clamp(min, max);
            }
        }

        0.0
    }

    fn is_learnable(&self) -> bool {
        true
    }

    fn get_parameters(&self) -> Vec<f64> {
        self.perturbation_weights.iter().copied().collect()
    }

    fn set_parameters(&mut self, params: &[f64]) -> PinnResult<()> {
        if params.len() != self.perturbation_weights.len() {
            return Err(PinnError::DimensionMismatch(format!(
                "Expected {} parameters, got {}",
                self.perturbation_weights.len(),
                params.len()
            )));
        }

        for (i, &p) in params.iter().enumerate() {
            let (row, col) = (i / self.base.n_layers(), i % self.base.n_layers());
            self.perturbation_weights[[row, col]] = p;
        }

        Ok(())
    }
}

/// Anisotropic conductivity tensor
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AnisotropicConductivity {
    /// Conductivity tensor [3, 3] at each layer
    tensors: Vec<Array2<f64>>,
    /// Base layered model for geometry
    geometry: LayeredConductivity,
}

impl AnisotropicConductivity {
    /// Create from isotropic layered model
    pub fn from_isotropic(geometry: LayeredConductivity) -> Self {
        let tensors = geometry
            .layers()
            .iter()
            .map(|l| {
                let sigma = l.conductivity;
                Array2::from_diag(&ndarray::arr1(&[sigma, sigma, sigma]))
            })
            .collect();

        Self { tensors, geometry }
    }

    /// Set anisotropic tensor for a layer
    pub fn set_tensor(&mut self, layer_idx: usize, tensor: Array2<f64>) -> PinnResult<()> {
        if layer_idx >= self.tensors.len() {
            return Err(PinnError::InvalidConfig(format!(
                "Layer index {} out of range",
                layer_idx
            )));
        }
        if tensor.shape() != [3, 3] {
            return Err(PinnError::DimensionMismatch("Tensor must be 3x3".into()));
        }
        self.tensors[layer_idx] = tensor;
        Ok(())
    }

    /// Get conductivity tensor at a point
    pub fn tensor_at(&self, point: &[f64; 3]) -> Array2<f64> {
        let dx = point[0] - self.geometry.center[0];
        let dy = point[1] - self.geometry.center[1];
        let dz = point[2] - self.geometry.center[2];
        let r = (dx * dx + dy * dy + dz * dz).sqrt();

        for (i, layer) in self.geometry.layers().iter().enumerate() {
            if layer.contains(r) {
                return self.tensors[i].clone();
            }
        }

        Array2::zeros((3, 3))
    }

    /// Compute σ∇Φ where σ is tensor
    pub fn apply_to_gradient(&self, point: &[f64; 3], gradient: &[f64; 3]) -> [f64; 3] {
        let tensor = self.tensor_at(point);
        let grad = ndarray::arr1(gradient);
        let result = tensor.dot(&grad);
        [result[0], result[1], result[2]]
    }
}

impl ConductivityModel for AnisotropicConductivity {
    fn conductivity_at(&self, point: &[f64; 3]) -> f64 {
        // Return trace / 3 as effective isotropic conductivity
        let tensor = self.tensor_at(point);
        (tensor[[0, 0]] + tensor[[1, 1]] + tensor[[2, 2]]) / 3.0
    }
}

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

    #[test]
    fn test_tissue_types() {
        assert!((TissueType::GrayMatter.default_conductivity() - 0.33).abs() < 1e-10);
        assert!((TissueType::Skull.default_conductivity() - 0.0042).abs() < 1e-10);
    }

    #[test]
    fn test_homogeneous() {
        let model = HomogeneousConductivity::new(0.33);
        assert!((model.conductivity_at(&[0.0, 0.0, 0.0]) - 0.33).abs() < 1e-10);
        assert!((model.conductivity_at(&[1.0, 2.0, 3.0]) - 0.33).abs() < 1e-10);
    }

    #[test]
    fn test_layered_conductivity() {
        let model = LayeredConductivity::three_layer(
            [0.0, 0.0, 0.0],
            0.08,  // brain radius
            0.007, // skull thickness
            0.006, // scalp thickness
        );

        assert_eq!(model.n_layers(), 3);

        // Inside brain
        let sigma = model.conductivity_at(&[0.0, 0.0, 0.05]);
        assert!((sigma - 0.33).abs() < 1e-10);

        // Inside skull
        let sigma = model.conductivity_at(&[0.0, 0.0, 0.082]);
        assert!((sigma - 0.0042).abs() < 1e-10);

        // Inside scalp
        let sigma = model.conductivity_at(&[0.0, 0.0, 0.09]);
        assert!((sigma - 0.43).abs() < 1e-10);
    }

    #[test]
    fn test_four_layer_model() {
        let model = LayeredConductivity::four_layer(
            [0.0, 0.0, 0.0],
            0.078, // brain
            0.002, // CSF
            0.007, // skull
            0.006, // scalp
        );

        assert_eq!(model.n_layers(), 4);

        // CSF layer
        let sigma = model.conductivity_at(&[0.0, 0.0, 0.079]);
        assert!((sigma - 1.79).abs() < 1e-10);
    }

    #[test]
    fn test_learnable_conductivity() {
        let base = LayeredConductivity::three_layer([0.0, 0.0, 0.0], 0.08, 0.007, 0.006);

        let model = LearnableConductivity::new(base, 5);

        // Should return values close to base model
        let sigma = model.conductivity_at(&[0.0, 0.0, 0.05]);
        assert!((sigma - 0.33).abs() < 0.05);

        // Has learnable parameters
        assert!(model.is_learnable());
        assert!(model.get_parameters().len() > 0);
    }

    #[test]
    fn test_anisotropic_conductivity() {
        let geometry = LayeredConductivity::three_layer([0.0, 0.0, 0.0], 0.08, 0.007, 0.006);

        let model = AnisotropicConductivity::from_isotropic(geometry);

        // Effective conductivity should equal isotropic value
        let sigma = model.conductivity_at(&[0.0, 0.0, 0.05]);
        assert!((sigma - 0.33).abs() < 1e-10);

        // Tensor should be diagonal
        let tensor = model.tensor_at(&[0.0, 0.0, 0.05]);
        assert!((tensor[[0, 0]] - 0.33).abs() < 1e-10);
        assert!(tensor[[0, 1]].abs() < 1e-10);
    }
}