//! Tissue conductivity models for FEM head modeling.
//!
//! Provides isotropic and anisotropic conductivity tensors
//! for different tissue types in the head.

use crate::error::{FemError, FemResult};
use crate::mesh::TissueLayer;
use nalgebra::{Matrix3, Vector3};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;

/// Conductivity tensor (3x3 symmetric matrix)
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ConductivityTensor {
    /// Tensor components (symmetric)
    pub tensor: Matrix3<f64>,
    /// Whether this is isotropic
    pub is_isotropic: bool,
}

impl ConductivityTensor {
    /// Create an isotropic conductivity tensor
    pub fn isotropic(sigma: f64) -> Self {
        Self {
            tensor: Matrix3::identity() * sigma,
            is_isotropic: true,
        }
    }

    /// Create an anisotropic conductivity tensor
    pub fn anisotropic(tensor: Matrix3<f64>) -> Self {
        Self {
            tensor,
            is_isotropic: false,
        }
    }

    /// Create from principal conductivities and eigenvectors
    pub fn from_principal(sigmas: &[f64; 3], eigenvectors: &Matrix3<f64>) -> Self {
        // σ = V * diag(σ1, σ2, σ3) * V^T
        let d = Matrix3::from_diagonal(&Vector3::new(sigmas[0], sigmas[1], sigmas[2]));
        let tensor = eigenvectors * d * eigenvectors.transpose();

        Self {
            tensor,
            is_isotropic: (sigmas[0] - sigmas[1]).abs() < 1e-10
                && (sigmas[1] - sigmas[2]).abs() < 1e-10,
        }
    }

    /// Get scalar conductivity (trace/3 for anisotropic)
    pub fn scalar(&self) -> f64 {
        (self.tensor[(0, 0)] + self.tensor[(1, 1)] + self.tensor[(2, 2)]) / 3.0
    }

    /// Get conductivity in a specific direction
    pub fn in_direction(&self, direction: &Vector3<f64>) -> f64 {
        let d = direction.normalize();
        d.dot(&(self.tensor * d))
    }

    /// Apply conductivity to electric field gradient
    pub fn apply(&self, gradient: &Vector3<f64>) -> Vector3<f64> {
        self.tensor * gradient
    }
}

impl Default for ConductivityTensor {
    fn default() -> Self {
        Self::isotropic(1.0)
    }
}

/// Anisotropy model type
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum AnisotropicModel {
    /// Isotropic (no anisotropy)
    Isotropic,
    /// Volume-based constraint
    VolumeConstraint,
    /// DTI-based from diffusion tensor
    DtiBased,
    /// Fixed ratio anisotropy
    FixedRatio,
}

/// Configuration for tissue conductivity
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TissueConfig {
    /// Tissue type
    pub tissue: TissueLayer,
    /// Isotropic conductivity (S/m)
    pub conductivity: f64,
    /// Anisotropy model
    pub anisotropy: AnisotropicModel,
    /// Anisotropy ratio (longitudinal/transverse) for non-DTI models
    pub anisotropy_ratio: f64,
}

impl TissueConfig {
    /// Create with isotropic conductivity
    pub fn isotropic(tissue: TissueLayer, conductivity: f64) -> Self {
        Self {
            tissue,
            conductivity,
            anisotropy: AnisotropicModel::Isotropic,
            anisotropy_ratio: 1.0,
        }
    }

    /// Create with anisotropic conductivity
    pub fn anisotropic(
        tissue: TissueLayer,
        conductivity: f64,
        model: AnisotropicModel,
        ratio: f64,
    ) -> Self {
        Self {
            tissue,
            conductivity,
            anisotropy: model,
            anisotropy_ratio: ratio,
        }
    }
}

/// Default conductivities for different tissue types (S/m)
pub fn default_conductivity(tissue: TissueLayer) -> f64 {
    match tissue {
        TissueLayer::Scalp => 0.43,
        TissueLayer::Skull => 0.0042,
        TissueLayer::Csf => 1.79,
        TissueLayer::GrayMatter => 0.33,
        TissueLayer::WhiteMatter => 0.14,
        TissueLayer::Air => 1e-12,
    }
}

/// Default anisotropy ratio for tissues
pub fn default_anisotropy_ratio(tissue: TissueLayer) -> f64 {
    match tissue {
        TissueLayer::Skull => 10.0,      // Radial/tangential
        TissueLayer::WhiteMatter => 9.0, // Along/perpendicular to fibers
        _ => 1.0,                        // Isotropic
    }
}

/// Tissue conductivity model for FEM
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TissueConductivity {
    /// Per-tissue configurations
    configs: HashMap<TissueLayer, TissueConfig>,
    /// Per-element conductivity tensors (if computed)
    element_tensors: Option<Vec<ConductivityTensor>>,
}

impl TissueConductivity {
    /// Create with default isotropic conductivities
    pub fn default_isotropic() -> Self {
        let mut configs = HashMap::new();

        for tissue in [
            TissueLayer::Scalp,
            TissueLayer::Skull,
            TissueLayer::Csf,
            TissueLayer::GrayMatter,
            TissueLayer::WhiteMatter,
        ] {
            configs.insert(
                tissue,
                TissueConfig::isotropic(tissue, default_conductivity(tissue)),
            );
        }

        Self {
            configs,
            element_tensors: None,
        }
    }

    /// Create with default anisotropic conductivities
    pub fn default_anisotropic() -> Self {
        let mut configs = HashMap::new();

        // Isotropic tissues
        for tissue in [
            TissueLayer::Scalp,
            TissueLayer::Csf,
            TissueLayer::GrayMatter,
        ] {
            configs.insert(
                tissue,
                TissueConfig::isotropic(tissue, default_conductivity(tissue)),
            );
        }

        // Anisotropic skull
        configs.insert(
            TissueLayer::Skull,
            TissueConfig::anisotropic(
                TissueLayer::Skull,
                default_conductivity(TissueLayer::Skull),
                AnisotropicModel::FixedRatio,
                default_anisotropy_ratio(TissueLayer::Skull),
            ),
        );

        // Anisotropic white matter
        configs.insert(
            TissueLayer::WhiteMatter,
            TissueConfig::anisotropic(
                TissueLayer::WhiteMatter,
                default_conductivity(TissueLayer::WhiteMatter),
                AnisotropicModel::VolumeConstraint,
                default_anisotropy_ratio(TissueLayer::WhiteMatter),
            ),
        );

        Self {
            configs,
            element_tensors: None,
        }
    }

    /// Set conductivity for a tissue
    pub fn set_conductivity(&mut self, tissue: TissueLayer, sigma: f64) {
        if let Some(config) = self.configs.get_mut(&tissue) {
            config.conductivity = sigma;
        } else {
            self.configs
                .insert(tissue, TissueConfig::isotropic(tissue, sigma));
        }
        self.element_tensors = None; // Invalidate cached tensors
    }

    /// Set anisotropy model for a tissue
    pub fn set_anisotropy(&mut self, tissue: TissueLayer, model: AnisotropicModel, ratio: f64) {
        if let Some(config) = self.configs.get_mut(&tissue) {
            config.anisotropy = model;
            config.anisotropy_ratio = ratio;
        }
        self.element_tensors = None;
    }

    /// Get configuration for a tissue
    pub fn get_config(&self, tissue: TissueLayer) -> Option<&TissueConfig> {
        self.configs.get(&tissue)
    }

    /// Get conductivity tensor for a tissue at a position
    pub fn tensor_at(
        &self,
        tissue: TissueLayer,
        position: &Vector3<f64>,
        normal: Option<&Vector3<f64>>,
    ) -> ConductivityTensor {
        let config = match self.configs.get(&tissue) {
            Some(c) => c,
            None => return ConductivityTensor::isotropic(default_conductivity(tissue)),
        };

        match config.anisotropy {
            AnisotropicModel::Isotropic => ConductivityTensor::isotropic(config.conductivity),
            AnisotropicModel::VolumeConstraint => {
                // Volume-preserving anisotropy
                self.volume_constraint_tensor(config, position, normal)
            }
            AnisotropicModel::FixedRatio => {
                // Fixed ratio along normal direction
                self.fixed_ratio_tensor(config, normal)
            }
            AnisotropicModel::DtiBased => {
                // Would need DTI data - fall back to isotropic
                ConductivityTensor::isotropic(config.conductivity)
            }
        }
    }

    /// Compute tensor with volume constraint
    fn volume_constraint_tensor(
        &self,
        config: &TissueConfig,
        position: &Vector3<f64>,
        _normal: Option<&Vector3<f64>>,
    ) -> ConductivityTensor {
        let sigma = config.conductivity;
        let ratio = config.anisotropy_ratio;

        // Get principal direction (radial from center for spherical)
        let r = position.norm();
        let radial = if r > 1e-10 {
            position / r
        } else {
            Vector3::new(0.0, 0.0, 1.0)
        };

        // Volume constraint: σ_l * σ_t^2 = σ_iso^3
        // With ratio = σ_l / σ_t
        let sigma_t = (sigma.powi(3) / ratio).powf(1.0 / 3.0);
        let sigma_l = sigma_t * ratio;

        // Build tensor: σ = σ_t * I + (σ_l - σ_t) * r ⊗ r
        let tensor =
            Matrix3::identity() * sigma_t + (sigma_l - sigma_t) * radial * radial.transpose();

        ConductivityTensor::anisotropic(tensor)
    }

    /// Compute tensor with fixed ratio along normal
    fn fixed_ratio_tensor(
        &self,
        config: &TissueConfig,
        normal: Option<&Vector3<f64>>,
    ) -> ConductivityTensor {
        let sigma = config.conductivity;
        let ratio = config.anisotropy_ratio;

        let default_normal = Vector3::new(0.0, 0.0, 1.0);
        let n = normal.unwrap_or(&default_normal);
        let n = n.normalize();

        // σ_radial = sigma * ratio, σ_tangential = sigma / sqrt(ratio)
        let sigma_n = sigma * ratio.sqrt();
        let sigma_t = sigma / ratio.sqrt();

        let tensor = Matrix3::identity() * sigma_t + (sigma_n - sigma_t) * n * n.transpose();

        ConductivityTensor::anisotropic(tensor)
    }

    /// Compute conductivity tensors for all elements
    pub fn compute_element_tensors(&mut self, mesh: &crate::mesh::HeadMesh) -> FemResult<()> {
        let n_elements = mesh.n_elements();
        let mut tensors = Vec::with_capacity(n_elements);

        for elem in &mesh.elements {
            let centroid = elem.centroid(&mesh.nodes);
            let normal = Some(centroid.normalize()); // Radial for spherical

            let tensor = self.tensor_at(elem.tissue, &centroid, normal.as_ref());
            tensors.push(tensor);
        }

        self.element_tensors = Some(tensors);
        Ok(())
    }

    /// Get pre-computed element tensor
    pub fn element_tensor(&self, element_idx: usize) -> Option<&ConductivityTensor> {
        self.element_tensors.as_ref()?.get(element_idx)
    }

    /// Get all element tensors
    pub fn element_tensors(&self) -> Option<&[ConductivityTensor]> {
        self.element_tensors.as_deref()
    }
}

impl Default for TissueConductivity {
    fn default() -> Self {
        Self::default_isotropic()
    }
}

/// DTI-based conductivity model
#[derive(Debug, Clone)]
pub struct DtiConductivity {
    /// Diffusion tensors per voxel
    pub diffusion_tensors: Vec<Matrix3<f64>>,
    /// Voxel positions
    pub positions: Vec<Vector3<f64>>,
    /// Conversion factor from diffusion to conductivity
    pub d2c_factor: f64,
    /// Minimum eigenvalue ratio
    pub min_ratio: f64,
}

impl DtiConductivity {
    /// Create from diffusion tensor data
    pub fn new(
        diffusion_tensors: Vec<Matrix3<f64>>,
        positions: Vec<Vector3<f64>>,
    ) -> FemResult<Self> {
        if diffusion_tensors.len() != positions.len() {
            return Err(FemError::DimensionMismatch(
                "Tensors and positions must have same length".into(),
            ));
        }

        Ok(Self {
            diffusion_tensors,
            positions,
            d2c_factor: 0.736, // Tuch 2001: σ = 0.736 * D
            min_ratio: 0.1,
        })
    }

    /// Get conductivity tensor at a position (nearest neighbor interpolation)
    pub fn tensor_at(&self, position: &Vector3<f64>) -> ConductivityTensor {
        if self.positions.is_empty() {
            return ConductivityTensor::isotropic(default_conductivity(TissueLayer::WhiteMatter));
        }

        // Find nearest DTI voxel
        let mut min_dist = f64::MAX;
        let mut nearest_idx = 0;

        for (i, pos) in self.positions.iter().enumerate() {
            let dist = (pos - position).norm_squared();
            if dist < min_dist {
                min_dist = dist;
                nearest_idx = i;
            }
        }

        // Convert diffusion to conductivity
        let d = &self.diffusion_tensors[nearest_idx];
        let sigma = d * self.d2c_factor;

        // Ensure positive definiteness
        let eigendecomp = sigma.symmetric_eigen();
        let mut eigenvalues = eigendecomp.eigenvalues;

        // Clamp eigenvalues
        let max_ev = eigenvalues.max();
        for ev in eigenvalues.iter_mut() {
            *ev = ev.max(max_ev * self.min_ratio);
        }

        let d_clamped = Matrix3::from_diagonal(&eigenvalues);
        let tensor = eigendecomp.eigenvectors * d_clamped * eigendecomp.eigenvectors.transpose();

        ConductivityTensor::anisotropic(tensor)
    }
}

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

    #[test]
    fn test_isotropic_tensor() {
        let tensor = ConductivityTensor::isotropic(0.33);
        assert!(tensor.is_isotropic);
        assert!((tensor.scalar() - 0.33).abs() < 1e-10);

        // Same in all directions
        let dir1 = Vector3::new(1.0, 0.0, 0.0);
        let dir2 = Vector3::new(0.0, 1.0, 0.0);
        assert!((tensor.in_direction(&dir1) - tensor.in_direction(&dir2)).abs() < 1e-10);
    }

    #[test]
    fn test_anisotropic_tensor() {
        let sigmas = [0.4, 0.2, 0.2];
        let eigenvectors = Matrix3::identity();
        let tensor = ConductivityTensor::from_principal(&sigmas, &eigenvectors);

        assert!(!tensor.is_isotropic);

        // Higher conductivity along x
        let dir_x = Vector3::new(1.0, 0.0, 0.0);
        let dir_y = Vector3::new(0.0, 1.0, 0.0);
        assert!(tensor.in_direction(&dir_x) > tensor.in_direction(&dir_y));
    }

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

    #[test]
    fn test_tissue_conductivity() {
        let conductivity = TissueConductivity::default_isotropic();

        let config = conductivity.get_config(TissueLayer::GrayMatter).unwrap();
        assert!((config.conductivity - 0.33).abs() < 1e-10);
        assert_eq!(config.anisotropy, AnisotropicModel::Isotropic);
    }

    #[test]
    fn test_anisotropic_conductivity() {
        let conductivity = TissueConductivity::default_anisotropic();

        let wm_config = conductivity.get_config(TissueLayer::WhiteMatter).unwrap();
        assert_eq!(wm_config.anisotropy, AnisotropicModel::VolumeConstraint);

        let skull_config = conductivity.get_config(TissueLayer::Skull).unwrap();
        assert_eq!(skull_config.anisotropy, AnisotropicModel::FixedRatio);
    }

    #[test]
    fn test_tensor_at() {
        let conductivity = TissueConductivity::default_anisotropic();

        let pos = Vector3::new(0.05, 0.0, 0.0);
        let tensor = conductivity.tensor_at(TissueLayer::WhiteMatter, &pos, None);

        assert!(!tensor.is_isotropic);
        assert!(tensor.scalar() > 0.0);
    }

    #[test]
    fn test_set_conductivity() {
        let mut conductivity = TissueConductivity::default_isotropic();

        conductivity.set_conductivity(TissueLayer::Skull, 0.01);
        let config = conductivity.get_config(TissueLayer::Skull).unwrap();
        assert!((config.conductivity - 0.01).abs() < 1e-10);
    }
}