//! Tetrahedral mesh generation for head models.
//!
//! Provides mesh generation from FreeSurfer surfaces and layered
//! spherical models for FEM forward solutions.

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

/// Tissue layer identifier
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum TissueLayer {
    /// Scalp/skin layer
    Scalp,
    /// Skull bone
    Skull,
    /// Cerebrospinal fluid
    Csf,
    /// Gray matter (cortex)
    GrayMatter,
    /// White matter
    WhiteMatter,
    /// Air (sinuses)
    Air,
}

impl TissueLayer {
    /// Get layer index for mesh material ID
    pub fn index(&self) -> usize {
        match self {
            TissueLayer::WhiteMatter => 0,
            TissueLayer::GrayMatter => 1,
            TissueLayer::Csf => 2,
            TissueLayer::Skull => 3,
            TissueLayer::Scalp => 4,
            TissueLayer::Air => 5,
        }
    }

    /// Get layer from index
    pub fn from_index(idx: usize) -> Option<Self> {
        match idx {
            0 => Some(TissueLayer::WhiteMatter),
            1 => Some(TissueLayer::GrayMatter),
            2 => Some(TissueLayer::Csf),
            3 => Some(TissueLayer::Skull),
            4 => Some(TissueLayer::Scalp),
            5 => Some(TissueLayer::Air),
            _ => None,
        }
    }
}

/// A node in the tetrahedral mesh
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MeshNode {
    /// Node ID
    pub id: usize,
    /// Position in 3D space (meters)
    pub position: Vector3<f64>,
    /// Whether this is a boundary node
    pub is_boundary: bool,
}

impl MeshNode {
    /// Create a new mesh node
    pub fn new(id: usize, x: f64, y: f64, z: f64) -> Self {
        Self {
            id,
            position: Vector3::new(x, y, z),
            is_boundary: false,
        }
    }
}

/// A tetrahedral element
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TetElement {
    /// Element ID
    pub id: usize,
    /// Node indices (4 nodes for linear tetrahedron)
    pub nodes: [usize; 4],
    /// Tissue layer this element belongs to
    pub tissue: TissueLayer,
    /// Element volume (m³)
    pub volume: f64,
}

impl TetElement {
    /// Create a new tetrahedral element
    pub fn new(id: usize, nodes: [usize; 4], tissue: TissueLayer) -> Self {
        Self {
            id,
            nodes,
            tissue,
            volume: 0.0,
        }
    }

    /// Compute element volume from node positions
    pub fn compute_volume(&mut self, mesh_nodes: &[MeshNode]) -> f64 {
        let p0 = &mesh_nodes[self.nodes[0]].position;
        let p1 = &mesh_nodes[self.nodes[1]].position;
        let p2 = &mesh_nodes[self.nodes[2]].position;
        let p3 = &mesh_nodes[self.nodes[3]].position;

        // Volume = |det([p1-p0, p2-p0, p3-p0])| / 6
        let v1 = p1 - p0;
        let v2 = p2 - p0;
        let v3 = p3 - p0;

        let det = v1.dot(&v2.cross(&v3));
        self.volume = det.abs() / 6.0;
        self.volume
    }

    /// Compute shape function gradients (constant for linear tet)
    pub fn shape_gradients(&self, mesh_nodes: &[MeshNode]) -> [Vector3<f64>; 4] {
        let p0 = &mesh_nodes[self.nodes[0]].position;
        let p1 = &mesh_nodes[self.nodes[1]].position;
        let p2 = &mesh_nodes[self.nodes[2]].position;
        let p3 = &mesh_nodes[self.nodes[3]].position;

        // Jacobian matrix
        let j = Matrix3::new(
            p1.x - p0.x,
            p2.x - p0.x,
            p3.x - p0.x,
            p1.y - p0.y,
            p2.y - p0.y,
            p3.y - p0.y,
            p1.z - p0.z,
            p2.z - p0.z,
            p3.z - p0.z,
        );

        let det = j.determinant();
        if det.abs() < 1e-15 {
            return [Vector3::zeros(); 4];
        }

        let j_inv = j.try_inverse().unwrap_or(Matrix3::identity());

        // Shape function gradients in reference coordinates
        // N0 = 1 - xi - eta - zeta, N1 = xi, N2 = eta, N3 = zeta
        let grad_ref = [
            Vector3::new(-1.0, -1.0, -1.0),
            Vector3::new(1.0, 0.0, 0.0),
            Vector3::new(0.0, 1.0, 0.0),
            Vector3::new(0.0, 0.0, 1.0),
        ];

        // Transform to physical coordinates
        [
            j_inv.transpose() * grad_ref[0],
            j_inv.transpose() * grad_ref[1],
            j_inv.transpose() * grad_ref[2],
            j_inv.transpose() * grad_ref[3],
        ]
    }

    /// Get element centroid
    pub fn centroid(&self, mesh_nodes: &[MeshNode]) -> Vector3<f64> {
        let mut c = Vector3::zeros();
        for &ni in &self.nodes {
            c += mesh_nodes[ni].position;
        }
        c / 4.0
    }
}

/// Mesh quality metrics
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MeshQuality {
    /// Minimum element volume
    pub min_volume: f64,
    /// Maximum element volume
    pub max_volume: f64,
    /// Average element volume
    pub avg_volume: f64,
    /// Minimum aspect ratio (1 = ideal)
    pub min_aspect_ratio: f64,
    /// Average aspect ratio
    pub avg_aspect_ratio: f64,
    /// Number of degenerate elements (volume < tolerance)
    pub n_degenerate: usize,
}

/// Head mesh for FEM forward modeling
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HeadMesh {
    /// Mesh nodes
    pub nodes: Vec<MeshNode>,
    /// Tetrahedral elements
    pub elements: Vec<TetElement>,
    /// Number of elements per tissue layer
    pub elements_per_layer: HashMap<TissueLayer, usize>,
    /// Mesh quality metrics
    pub quality: Option<MeshQuality>,
}

impl HeadMesh {
    /// Create a new empty head mesh
    pub fn new() -> Self {
        Self {
            nodes: Vec::new(),
            elements: Vec::new(),
            elements_per_layer: HashMap::new(),
            quality: None,
        }
    }

    /// Create a layered spherical head mesh
    ///
    /// # Arguments
    /// * `radii` - Radii for each layer boundary (inner to outer)
    /// * `n_radial` - Number of radial divisions per layer
    /// * `n_angular` - Number of angular divisions (icosahedral refinement level)
    pub fn spherical(
        radii: &[(TissueLayer, f64)],
        n_radial: usize,
        n_angular: usize,
    ) -> FemResult<Self> {
        let mut mesh = HeadMesh::new();

        if radii.is_empty() {
            return Err(FemError::InvalidMesh("No layer radii specified".into()));
        }

        // Generate nodes on concentric spherical shells
        let n_shells = radii.len() * n_radial + 1;
        let mut shell_radii = Vec::with_capacity(n_shells);

        // Inner point at center
        shell_radii.push(0.0);

        for (i, (_, outer_r)) in radii.iter().enumerate() {
            let inner_r = if i == 0 { 0.0 } else { radii[i - 1].1 };
            for j in 1..=n_radial {
                let t = j as f64 / n_radial as f64;
                let r = inner_r + t * (outer_r - inner_r);
                shell_radii.push(r);
            }
        }

        // Generate icosahedral points on each shell
        let ico_points = generate_icosahedral_points(n_angular);
        let n_points_per_shell = ico_points.len();

        // Center node
        mesh.nodes.push(MeshNode::new(0, 0.0, 0.0, 0.0));

        // Nodes on shells
        for (shell_idx, &r) in shell_radii.iter().enumerate().skip(1) {
            for (pt_idx, pt) in ico_points.iter().enumerate() {
                let node_id = 1 + (shell_idx - 1) * n_points_per_shell + pt_idx;
                let pos = pt * r;
                let mut node = MeshNode::new(node_id, pos.x, pos.y, pos.z);
                node.is_boundary = shell_idx == shell_radii.len() - 1;
                mesh.nodes.push(node);
            }
        }

        // Generate tetrahedral elements
        // Connect center to first shell
        let first_shell_start = 1;
        for tri in generate_icosahedral_triangles(n_angular) {
            let n0 = 0; // center
            let n1 = first_shell_start + tri[0];
            let n2 = first_shell_start + tri[1];
            let n3 = first_shell_start + tri[2];

            let tissue = radii[0].0;
            let elem_id = mesh.elements.len();
            mesh.elements
                .push(TetElement::new(elem_id, [n0, n1, n2, n3], tissue));
        }

        // Connect adjacent shells with prism-split tetrahedra
        for shell in 1..shell_radii.len() - 1 {
            let inner_start = 1 + (shell - 1) * n_points_per_shell;
            let outer_start = inner_start + n_points_per_shell;

            // Determine tissue layer
            let layer_idx = shell / n_radial;
            let tissue = radii[layer_idx.min(radii.len() - 1)].0;

            for tri in generate_icosahedral_triangles(n_angular) {
                // Inner and outer triangle nodes
                let i0 = inner_start + tri[0];
                let i1 = inner_start + tri[1];
                let i2 = inner_start + tri[2];
                let o0 = outer_start + tri[0];
                let o1 = outer_start + tri[1];
                let o2 = outer_start + tri[2];

                // Split prism into 3 tetrahedra
                let elem_id = mesh.elements.len();
                mesh.elements
                    .push(TetElement::new(elem_id, [i0, i1, i2, o0], tissue));
                mesh.elements
                    .push(TetElement::new(elem_id + 1, [i1, i2, o0, o1], tissue));
                mesh.elements
                    .push(TetElement::new(elem_id + 2, [i2, o0, o1, o2], tissue));
            }
        }

        // Compute element volumes and count per layer
        for elem in &mut mesh.elements {
            elem.compute_volume(&mesh.nodes);
            *mesh.elements_per_layer.entry(elem.tissue).or_insert(0) += 1;
        }

        mesh.compute_quality();

        Ok(mesh)
    }

    /// Create a simple 3-layer spherical head
    pub fn three_layer_sphere(
        brain_radius: f64,
        skull_thickness: f64,
        scalp_thickness: f64,
        n_radial: usize,
        n_angular: usize,
    ) -> FemResult<Self> {
        let radii = vec![
            (TissueLayer::GrayMatter, brain_radius),
            (TissueLayer::Skull, brain_radius + skull_thickness),
            (
                TissueLayer::Scalp,
                brain_radius + skull_thickness + scalp_thickness,
            ),
        ];

        Self::spherical(&radii, n_radial, n_angular)
    }

    /// Create a 5-layer spherical head (with CSF and white matter)
    pub fn five_layer_sphere(
        white_radius: f64,
        gray_thickness: f64,
        csf_thickness: f64,
        skull_thickness: f64,
        scalp_thickness: f64,
        n_radial: usize,
        n_angular: usize,
    ) -> FemResult<Self> {
        let gray_r = white_radius + gray_thickness;
        let csf_r = gray_r + csf_thickness;
        let skull_r = csf_r + skull_thickness;
        let scalp_r = skull_r + scalp_thickness;

        let radii = vec![
            (TissueLayer::WhiteMatter, white_radius),
            (TissueLayer::GrayMatter, gray_r),
            (TissueLayer::Csf, csf_r),
            (TissueLayer::Skull, skull_r),
            (TissueLayer::Scalp, scalp_r),
        ];

        Self::spherical(&radii, n_radial, n_angular)
    }

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

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

    /// Compute mesh quality metrics
    pub fn compute_quality(&mut self) {
        if self.elements.is_empty() {
            return;
        }

        let mut min_vol = f64::MAX;
        let mut max_vol: f64 = 0.0;
        let mut sum_vol: f64 = 0.0;
        let mut n_degenerate = 0;
        let mut sum_aspect = 0.0;
        let mut min_aspect = f64::MAX;
        let vol_tol = 1e-20;

        for elem in &self.elements {
            let vol = elem.volume;
            min_vol = min_vol.min(vol);
            max_vol = max_vol.max(vol);
            sum_vol += vol;

            if vol < vol_tol {
                n_degenerate += 1;
            }

            // Compute aspect ratio (edge length ratio)
            let aspect = self.element_aspect_ratio(elem);
            sum_aspect += aspect;
            min_aspect = min_aspect.min(aspect);
        }

        let n = self.elements.len() as f64;
        self.quality = Some(MeshQuality {
            min_volume: min_vol,
            max_volume: max_vol,
            avg_volume: sum_vol / n,
            min_aspect_ratio: min_aspect,
            avg_aspect_ratio: sum_aspect / n,
            n_degenerate,
        });
    }

    /// Compute aspect ratio for an element (shortest edge / longest edge)
    fn element_aspect_ratio(&self, elem: &TetElement) -> f64 {
        let edges = [(0, 1), (0, 2), (0, 3), (1, 2), (1, 3), (2, 3)];

        let mut min_len = f64::MAX;
        let mut max_len: f64 = 0.0;

        for (i, j) in edges {
            let p1 = &self.nodes[elem.nodes[i]].position;
            let p2 = &self.nodes[elem.nodes[j]].position;
            let len = (p1 - p2).norm();
            min_len = min_len.min(len);
            max_len = max_len.max(len);
        }

        if max_len > 1e-15 {
            min_len / max_len
        } else {
            0.0
        }
    }

    /// Get total mesh volume
    pub fn total_volume(&self) -> f64 {
        self.elements.iter().map(|e| e.volume).sum()
    }

    /// Get volume per tissue layer
    pub fn volume_per_layer(&self) -> HashMap<TissueLayer, f64> {
        let mut volumes = HashMap::new();
        for elem in &self.elements {
            *volumes.entry(elem.tissue).or_insert(0.0) += elem.volume;
        }
        volumes
    }

    /// Find element containing a point
    pub fn find_element(&self, point: &Vector3<f64>) -> Option<usize> {
        for (i, elem) in self.elements.iter().enumerate() {
            if self.point_in_element(elem, point) {
                return Some(i);
            }
        }
        None
    }

    /// Check if a point is inside an element using barycentric coordinates
    fn point_in_element(&self, elem: &TetElement, point: &Vector3<f64>) -> bool {
        let p0 = &self.nodes[elem.nodes[0]].position;
        let p1 = &self.nodes[elem.nodes[1]].position;
        let p2 = &self.nodes[elem.nodes[2]].position;
        let p3 = &self.nodes[elem.nodes[3]].position;

        // Compute barycentric coordinates
        let v0 = p1 - p0;
        let v1 = p2 - p0;
        let v2 = p3 - p0;
        let vp = point - p0;

        let d00 = v0.dot(&v0);
        let d01 = v0.dot(&v1);
        let d02 = v0.dot(&v2);
        let d11 = v1.dot(&v1);
        let d12 = v1.dot(&v2);
        let d22 = v2.dot(&v2);
        let dp0 = vp.dot(&v0);
        let dp1 = vp.dot(&v1);
        let dp2 = vp.dot(&v2);

        let det = d00 * (d11 * d22 - d12 * d12) - d01 * (d01 * d22 - d12 * d02)
            + d02 * (d01 * d12 - d11 * d02);

        if det.abs() < 1e-15 {
            return false;
        }

        let inv_det = 1.0 / det;

        // Compute barycentric coordinates
        let l1 = inv_det
            * (dp0 * (d11 * d22 - d12 * d12)
                + dp1 * (d02 * d12 - d01 * d22)
                + dp2 * (d01 * d12 - d02 * d11));
        let l2 = inv_det
            * (dp0 * (d12 * d02 - d01 * d22)
                + dp1 * (d00 * d22 - d02 * d02)
                + dp2 * (d01 * d02 - d00 * d12));
        let l3 = inv_det
            * (dp0 * (d01 * d12 - d11 * d02)
                + dp1 * (d01 * d02 - d00 * d12)
                + dp2 * (d00 * d11 - d01 * d01));
        let l0 = 1.0 - l1 - l2 - l3;

        let eps = -1e-10;
        l0 >= eps && l1 >= eps && l2 >= eps && l3 >= eps
    }

    /// Get node positions as array
    pub fn node_positions(&self) -> Array2<f64> {
        let n = self.nodes.len();
        let mut pos = Array2::zeros((n, 3));
        for (i, node) in self.nodes.iter().enumerate() {
            pos[[i, 0]] = node.position.x;
            pos[[i, 1]] = node.position.y;
            pos[[i, 2]] = node.position.z;
        }
        pos
    }

    /// Get element connectivity as array
    pub fn element_connectivity(&self) -> Array2<usize> {
        let n = self.elements.len();
        let mut conn = Array2::zeros((n, 4));
        for (i, elem) in self.elements.iter().enumerate() {
            for j in 0..4 {
                conn[[i, j]] = elem.nodes[j];
            }
        }
        conn
    }
}

impl Default for HeadMesh {
    fn default() -> Self {
        Self::new()
    }
}

/// Generate icosahedral sphere mesh (vertices and faces) at given refinement level
fn generate_icosahedral_mesh(refinement: usize) -> (Vec<Vector3<f64>>, Vec<[usize; 3]>) {
    // Base icosahedron vertices
    let phi = f64::midpoint(1.0, 5.0_f64.sqrt());
    let scale = 1.0 / (1.0 + phi * phi).sqrt();

    let mut vertices = vec![
        Vector3::new(0.0, 1.0, phi) * scale,
        Vector3::new(0.0, -1.0, phi) * scale,
        Vector3::new(0.0, 1.0, -phi) * scale,
        Vector3::new(0.0, -1.0, -phi) * scale,
        Vector3::new(1.0, phi, 0.0) * scale,
        Vector3::new(-1.0, phi, 0.0) * scale,
        Vector3::new(1.0, -phi, 0.0) * scale,
        Vector3::new(-1.0, -phi, 0.0) * scale,
        Vector3::new(phi, 0.0, 1.0) * scale,
        Vector3::new(-phi, 0.0, 1.0) * scale,
        Vector3::new(phi, 0.0, -1.0) * scale,
        Vector3::new(-phi, 0.0, -1.0) * scale,
    ];

    let mut faces = vec![
        [0, 1, 8],
        [0, 8, 4],
        [0, 4, 5],
        [0, 5, 9],
        [0, 9, 1],
        [1, 6, 8],
        [8, 6, 10],
        [8, 10, 4],
        [4, 10, 2],
        [4, 2, 5],
        [5, 2, 11],
        [5, 11, 9],
        [9, 11, 7],
        [9, 7, 1],
        [1, 7, 6],
        [3, 6, 7],
        [3, 7, 11],
        [3, 11, 2],
        [3, 2, 10],
        [3, 10, 6],
    ];

    // Refine by subdividing triangles
    for _ in 0..refinement {
        let mut new_faces = Vec::with_capacity(faces.len() * 4);
        let mut edge_midpoints: HashMap<(usize, usize), usize> = HashMap::new();

        for face in &faces {
            let mut mid = [0usize; 3];

            for i in 0..3 {
                let (a, b) = (face[i], face[(i + 1) % 3]);
                let key = if a < b { (a, b) } else { (b, a) };

                mid[i] = *edge_midpoints.entry(key).or_insert_with(|| {
                    let midpoint = (vertices[a] + vertices[b]).normalize();
                    vertices.push(midpoint);
                    vertices.len() - 1
                });
            }

            // Create 4 new triangles
            new_faces.push([face[0], mid[0], mid[2]]);
            new_faces.push([face[1], mid[1], mid[0]]);
            new_faces.push([face[2], mid[2], mid[1]]);
            new_faces.push([mid[0], mid[1], mid[2]]);
        }

        faces = new_faces;
    }

    (vertices, faces)
}

/// Generate points on an icosahedral sphere at given refinement level
fn generate_icosahedral_points(refinement: usize) -> Vec<Vector3<f64>> {
    let (vertices, _) = generate_icosahedral_mesh(refinement);
    vertices
}

/// Generate icosahedral triangle indices at given refinement level
fn generate_icosahedral_triangles(refinement: usize) -> Vec<[usize; 3]> {
    let (_, faces) = generate_icosahedral_mesh(refinement);
    faces
}

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

    #[test]
    fn test_tissue_layer_index() {
        assert_eq!(TissueLayer::WhiteMatter.index(), 0);
        assert_eq!(TissueLayer::Scalp.index(), 4);
        assert_eq!(TissueLayer::from_index(0), Some(TissueLayer::WhiteMatter));
    }

    #[test]
    fn test_mesh_node() {
        let node = MeshNode::new(0, 1.0, 2.0, 3.0);
        assert_eq!(node.id, 0);
        assert!((node.position.x - 1.0).abs() < 1e-10);
    }

    #[test]
    fn test_tet_element_volume() {
        let nodes = vec![
            MeshNode::new(0, 0.0, 0.0, 0.0),
            MeshNode::new(1, 1.0, 0.0, 0.0),
            MeshNode::new(2, 0.0, 1.0, 0.0),
            MeshNode::new(3, 0.0, 0.0, 1.0),
        ];

        let mut elem = TetElement::new(0, [0, 1, 2, 3], TissueLayer::GrayMatter);
        let vol = elem.compute_volume(&nodes);

        // Volume of unit tetrahedron = 1/6
        assert!((vol - 1.0 / 6.0).abs() < 1e-10);
    }

    #[test]
    fn test_icosahedral_points() {
        let points = generate_icosahedral_points(0);
        assert_eq!(points.len(), 12); // Base icosahedron

        let points = generate_icosahedral_points(1);
        assert!(points.len() > 12); // Refined

        // All points should be on unit sphere
        for p in &points {
            assert!((p.norm() - 1.0).abs() < 1e-10);
        }
    }

    #[test]
    fn test_three_layer_sphere() {
        let mesh = HeadMesh::three_layer_sphere(
            0.08,  // brain
            0.007, // skull
            0.006, // scalp
            2,     // radial divisions
            1,     // angular refinement
        )
        .unwrap();

        assert!(mesh.n_nodes() > 0);
        assert!(mesh.n_elements() > 0);

        // Should have elements in each layer
        assert!(
            mesh.elements_per_layer
                .contains_key(&TissueLayer::GrayMatter)
        );
        assert!(mesh.elements_per_layer.contains_key(&TissueLayer::Skull));
        assert!(mesh.elements_per_layer.contains_key(&TissueLayer::Scalp));
    }

    #[test]
    fn test_five_layer_sphere() {
        let mesh = HeadMesh::five_layer_sphere(
            0.06,  // white matter
            0.015, // gray matter
            0.002, // CSF
            0.007, // skull
            0.006, // scalp
            2, 1,
        )
        .unwrap();

        assert!(mesh.n_nodes() > 0);
        assert!(mesh.n_elements() > 0);
        assert!(mesh.elements_per_layer.len() == 5);
    }

    #[test]
    fn test_mesh_quality() {
        let mesh = HeadMesh::three_layer_sphere(0.08, 0.007, 0.006, 2, 1).unwrap();

        assert!(mesh.quality.is_some());
        let q = mesh.quality.as_ref().unwrap();
        assert!(q.min_volume > 0.0);
        assert!(q.avg_volume > 0.0);
        assert!(q.avg_aspect_ratio > 0.0);
        assert!(q.avg_aspect_ratio <= 1.0);
    }

    #[test]
    fn test_find_element() {
        let mesh = HeadMesh::three_layer_sphere(0.08, 0.007, 0.006, 2, 1).unwrap();

        // Point inside mesh
        let center = Vector3::new(0.0, 0.0, 0.0);
        let elem_idx = mesh.find_element(&center);
        assert!(elem_idx.is_some());

        // Point outside mesh
        let outside = Vector3::new(1.0, 0.0, 0.0);
        let elem_idx = mesh.find_element(&outside);
        assert!(elem_idx.is_none());
    }
}