//! Boundary Element Method (BEM) forward model.
//!
//! BEM provides accurate EEG forward modeling by accounting for the
//! realistic geometry of the head's conducting layers.
//!
//! ## Head Model
//!
//! The BEM model uses three surfaces:
//! - **Inner skull**: Boundary between brain and skull
//! - **Outer skull**: Boundary between skull and scalp
//! - **Scalp**: Outer head surface
//!
//! ## Algorithm
//!
//! Uses the symmetric BEM formulation:
//! 1. Compute BEM matrix from surface geometries
//! 2. Solve linear system for surface potentials
//! 3. Compute lead field from surface potentials to electrodes
//!
//! ## Reference
//!
//! Gramfort, A., et al. (2010). OpenMEEG: opensource software for quasistatic
//! bioelectromagnetics. Biomedical engineering online, 9(1), 45.

use crate::gain::GainMatrix;
use crate::sensors::{SensorArray, SensorType};
use crate::source_space::SourceSpace;
use crate::{ForwardError, ForwardResult, Orientation, Position, cross, dot, norm, normalize};
use nalgebra::{DMatrix, DVector, Vector3};
use std::f64::consts::PI;

/// A triangular mesh surface for BEM
#[derive(Debug, Clone)]
pub struct BemSurface {
    /// Vertex positions
    vertices: Vec<Position>,
    /// Triangle indices (each triangle is 3 vertex indices)
    triangles: Vec<[usize; 3]>,
    /// Surface name (e.g., "inner_skull", "outer_skull", "scalp")
    name: String,
    /// Precomputed triangle centroids
    centroids: Vec<Position>,
    /// Precomputed triangle normals
    normals: Vec<Orientation>,
    /// Precomputed triangle areas
    areas: Vec<f64>,
}

impl BemSurface {
    /// Create a BEM surface from vertices and triangles
    pub fn new(
        vertices: Vec<[f64; 3]>,
        triangles: Vec<[usize; 3]>,
        name: &str,
    ) -> ForwardResult<Self> {
        if vertices.is_empty() {
            return Err(ForwardError::InvalidGeometry(
                "Empty vertex list".to_string(),
            ));
        }

        if triangles.is_empty() {
            return Err(ForwardError::InvalidGeometry(
                "Empty triangle list".to_string(),
            ));
        }

        let vertices: Vec<Position> = vertices
            .into_iter()
            .map(|v| Vector3::new(v[0], v[1], v[2]))
            .collect();

        // Validate triangle indices
        let n_vertices = vertices.len();
        for tri in &triangles {
            for &idx in tri {
                if idx >= n_vertices {
                    return Err(ForwardError::InvalidGeometry(format!(
                        "Triangle index {} out of bounds (n_vertices={})",
                        idx, n_vertices
                    )));
                }
            }
        }

        // Precompute triangle properties
        let mut centroids = Vec::with_capacity(triangles.len());
        let mut normals = Vec::with_capacity(triangles.len());
        let mut areas = Vec::with_capacity(triangles.len());

        for tri in &triangles {
            let v0 = &vertices[tri[0]];
            let v1 = &vertices[tri[1]];
            let v2 = &vertices[tri[2]];

            // Centroid
            let centroid = (v0 + v1 + v2) / 3.0;
            centroids.push(centroid);

            // Normal (cross product of two edges)
            let e1 = v1 - v0;
            let e2 = v2 - v0;
            let n = cross(&e1, &e2);
            let area = norm(&n) / 2.0;
            let normal = if area > 1e-15 {
                normalize(&n)
            } else {
                Vector3::zeros()
            };

            normals.push(normal);
            areas.push(area);
        }

        Ok(Self {
            vertices,
            triangles,
            name: name.to_string(),
            centroids,
            normals,
            areas,
        })
    }

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

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

    /// Get surface name
    pub fn name(&self) -> &str {
        &self.name
    }

    /// Get triangle centroids
    pub fn centroids(&self) -> &[Position] {
        &self.centroids
    }

    /// Get triangle normals
    pub fn normals(&self) -> &[Orientation] {
        &self.normals
    }

    /// Get triangle areas
    pub fn areas(&self) -> &[f64] {
        &self.areas
    }

    /// Get vertices
    pub fn vertices(&self) -> &[Position] {
        &self.vertices
    }

    /// Create a spherical BEM surface for testing
    pub fn sphere(
        center: [f64; 3],
        radius: f64,
        n_subdivisions: usize,
        name: &str,
    ) -> ForwardResult<Self> {
        // Create icosphere
        let (vertices, triangles) = Self::create_icosphere(center, radius, n_subdivisions);
        Self::new(vertices, triangles, name)
    }

    /// Create an icosphere (subdivided icosahedron)
    fn create_icosphere(
        center: [f64; 3],
        radius: f64,
        subdivisions: usize,
    ) -> (Vec<[f64; 3]>, Vec<[usize; 3]>) {
        // Golden ratio
        let phi = f64::midpoint(1.0, 5.0_f64.sqrt());
        let len = (1.0 + phi * phi).sqrt();

        // Icosahedron vertices
        let mut vertices: Vec<[f64; 3]> = vec![
            [-1.0 / len, phi / len, 0.0],
            [1.0 / len, phi / len, 0.0],
            [-1.0 / len, -phi / len, 0.0],
            [1.0 / len, -phi / len, 0.0],
            [0.0, -1.0 / len, phi / len],
            [0.0, 1.0 / len, phi / len],
            [0.0, -1.0 / len, -phi / len],
            [0.0, 1.0 / len, -phi / len],
            [phi / len, 0.0, -1.0 / len],
            [phi / len, 0.0, 1.0 / len],
            [-phi / len, 0.0, -1.0 / len],
            [-phi / len, 0.0, 1.0 / len],
        ];

        // Scale to radius and translate to center
        for v in &mut vertices {
            let n = (v[0] * v[0] + v[1] * v[1] + v[2] * v[2]).sqrt();
            v[0] = center[0] + radius * v[0] / n;
            v[1] = center[1] + radius * v[1] / n;
            v[2] = center[2] + radius * v[2] / n;
        }

        // Icosahedron faces
        let mut triangles: Vec<[usize; 3]> = vec![
            [0, 11, 5],
            [0, 5, 1],
            [0, 1, 7],
            [0, 7, 10],
            [0, 10, 11],
            [1, 5, 9],
            [5, 11, 4],
            [11, 10, 2],
            [10, 7, 6],
            [7, 1, 8],
            [3, 9, 4],
            [3, 4, 2],
            [3, 2, 6],
            [3, 6, 8],
            [3, 8, 9],
            [4, 9, 5],
            [2, 4, 11],
            [6, 2, 10],
            [8, 6, 7],
            [9, 8, 1],
        ];

        // Subdivide
        for _ in 0..subdivisions {
            let mut new_triangles = Vec::new();
            let mut midpoint_cache: std::collections::HashMap<(usize, usize), usize> =
                std::collections::HashMap::new();

            for tri in &triangles {
                // Get midpoints
                let a = Self::get_midpoint(
                    &mut vertices,
                    &mut midpoint_cache,
                    tri[0],
                    tri[1],
                    center,
                    radius,
                );
                let b = Self::get_midpoint(
                    &mut vertices,
                    &mut midpoint_cache,
                    tri[1],
                    tri[2],
                    center,
                    radius,
                );
                let c = Self::get_midpoint(
                    &mut vertices,
                    &mut midpoint_cache,
                    tri[2],
                    tri[0],
                    center,
                    radius,
                );

                // Create 4 new triangles
                new_triangles.push([tri[0], a, c]);
                new_triangles.push([tri[1], b, a]);
                new_triangles.push([tri[2], c, b]);
                new_triangles.push([a, b, c]);
            }

            triangles = new_triangles;
        }

        (vertices, triangles)
    }

    /// Get or create midpoint vertex
    fn get_midpoint(
        vertices: &mut Vec<[f64; 3]>,
        cache: &mut std::collections::HashMap<(usize, usize), usize>,
        i1: usize,
        i2: usize,
        center: [f64; 3],
        radius: f64,
    ) -> usize {
        let key = if i1 < i2 { (i1, i2) } else { (i2, i1) };

        if let Some(&idx) = cache.get(&key) {
            return idx;
        }

        let v1 = &vertices[i1];
        let v2 = &vertices[i2];

        // Midpoint
        let mut mid = [
            f64::midpoint(v1[0], v2[0]),
            f64::midpoint(v1[1], v2[1]),
            f64::midpoint(v1[2], v2[2]),
        ];

        // Project to sphere
        let dx = mid[0] - center[0];
        let dy = mid[1] - center[1];
        let dz = mid[2] - center[2];
        let len = (dx * dx + dy * dy + dz * dz).sqrt();

        mid[0] = center[0] + radius * dx / len;
        mid[1] = center[1] + radius * dy / len;
        mid[2] = center[2] + radius * dz / len;

        let idx = vertices.len();
        vertices.push(mid);
        cache.insert(key, idx);
        idx
    }
}

/// Configuration for BEM head model
#[derive(Debug, Clone)]
pub struct BemConfig {
    /// Conductivity values [brain, skull, scalp] in S/m
    pub conductivities: [f64; 3],
    /// Use symmetric BEM formulation
    pub symmetric: bool,
    /// Integration order for computing BEM coefficients
    pub integration_order: usize,
}

impl Default for BemConfig {
    fn default() -> Self {
        Self {
            // Standard conductivity values (S/m)
            conductivities: [0.33, 0.0042, 0.33], // brain, skull, scalp
            symmetric: true,
            integration_order: 3,
        }
    }
}

impl BemConfig {
    /// Create config with custom conductivities
    pub fn with_conductivities(brain: f64, skull: f64, scalp: f64) -> Self {
        Self {
            conductivities: [brain, skull, scalp],
            ..Self::default()
        }
    }
}

/// BEM forward model for EEG
#[derive(Debug)]
pub struct BemModel {
    /// Surfaces (ordered: inner_skull, outer_skull, scalp)
    surfaces: Vec<BemSurface>,
    /// Configuration
    config: BemConfig,
    /// BEM system matrix (precomputed)
    bem_matrix: Option<DMatrix<f64>>,
    /// Total number of collocation points
    n_collocation: usize,
}

impl BemModel {
    /// Create a BEM model from three surfaces
    ///
    /// # Arguments
    /// * `inner_skull` - Brain-skull boundary surface
    /// * `outer_skull` - Skull-scalp boundary surface
    /// * `scalp` - Outer head surface
    /// * `config` - Optional configuration
    pub fn new(
        inner_skull: BemSurface,
        outer_skull: BemSurface,
        scalp: BemSurface,
        config: Option<BemConfig>,
    ) -> ForwardResult<Self> {
        let config = config.unwrap_or_default();

        let n_collocation =
            inner_skull.n_triangles() + outer_skull.n_triangles() + scalp.n_triangles();

        Ok(Self {
            surfaces: vec![inner_skull, outer_skull, scalp],
            config,
            bem_matrix: None,
            n_collocation,
        })
    }

    /// Create a three-shell spherical BEM model for testing
    pub fn three_shell_sphere(
        center: [f64; 3],
        radii: [f64; 3], // [inner_skull, outer_skull, scalp]
        n_triangles: usize,
        config: Option<BemConfig>,
    ) -> ForwardResult<Self> {
        // Determine subdivision level to get approximately n_triangles
        let subdivisions = ((n_triangles as f64 / 20.0).log2() / 2.0).ceil() as usize;

        let inner_skull = BemSurface::sphere(center, radii[0], subdivisions, "inner_skull")?;
        let outer_skull = BemSurface::sphere(center, radii[1], subdivisions, "outer_skull")?;
        let scalp = BemSurface::sphere(center, radii[2], subdivisions, "scalp")?;

        Self::new(inner_skull, outer_skull, scalp, config)
    }

    /// Compute the BEM system matrix
    ///
    /// This precomputes the matrix needed for the forward solution.
    /// Call this before computing gain matrices.
    pub fn make_bem_matrix(&mut self) -> ForwardResult<()> {
        let n = self.n_collocation;
        let mut matrix = DMatrix::zeros(n, n);

        let sigma = &self.config.conductivities;

        // Compute sigma ratios
        // sigma_minus / (sigma_plus + sigma_minus) for each interface
        let sigma_ratios = [
            sigma[0] / (sigma[0] + sigma[1]), // inner skull: brain-skull
            sigma[1] / (sigma[1] + sigma[2]), // outer skull: skull-scalp
            1.0,                              // scalp: scalp-air (sigma_air ≈ 0)
        ];

        // Fill BEM matrix using linear collocation
        let mut row_offset = 0;
        for (i, surf_i) in self.surfaces.iter().enumerate() {
            let mut col_offset = 0;

            for (j, surf_j) in self.surfaces.iter().enumerate() {
                // Compute interaction block
                let block = self.compute_bem_block(surf_i, surf_j, i == j);

                // Scale by conductivity ratio
                let scale = if i == j {
                    0.5 - sigma_ratios[i]
                } else if j < i {
                    -sigma_ratios[j]
                } else {
                    sigma_ratios[i] - 1.0
                };

                // Copy block to matrix
                for (bi, _ti) in surf_i.centroids().iter().enumerate() {
                    for (bj, _tj) in surf_j.centroids().iter().enumerate() {
                        matrix[(row_offset + bi, col_offset + bj)] = block[(bi, bj)] * scale;
                    }
                }

                // Add identity on diagonal for same surface
                if i == j {
                    for bi in 0..surf_i.n_triangles() {
                        matrix[(row_offset + bi, col_offset + bi)] += 0.5;
                    }
                }

                col_offset += surf_j.n_triangles();
            }

            row_offset += surf_i.n_triangles();
        }

        self.bem_matrix = Some(matrix);
        Ok(())
    }

    /// Compute the gain matrix for EEG electrodes
    pub fn compute_gain(
        &self,
        sources: &SourceSpace,
        sensors: &SensorArray,
    ) -> ForwardResult<GainMatrix> {
        let bem_matrix = self.bem_matrix.as_ref().ok_or_else(|| {
            ForwardError::ComputationError(
                "BEM matrix not computed. Call make_bem_matrix() first".to_string(),
            )
        })?;

        // Invert BEM matrix
        let bem_inv = Self::invert_matrix(bem_matrix)?;

        let n_sensors = sensors.len();
        let n_sources = sources.len();
        let free_ori = !sources.is_fixed_orientation();

        let n_columns = if free_ori { 3 * n_sources } else { n_sources };
        let mut gain_data = vec![vec![0.0; n_columns]; n_sensors];

        // Get EEG sensor positions
        let sensor_positions: Vec<&Position> = sensors
            .iter()
            .filter(|s| matches!(s.sensor_type(), SensorType::Eeg))
            .map(super::sensors::Sensor::position)
            .collect();

        // For each source, compute the lead field
        for (src_idx, source) in sources.iter().enumerate() {
            let dipole_pos = source.position();

            if free_ori {
                // Compute for each orientation
                for ori_idx in 0..3 {
                    let dipole_ori = match ori_idx {
                        0 => Vector3::new(1.0, 0.0, 0.0),
                        1 => Vector3::new(0.0, 1.0, 0.0),
                        _ => Vector3::new(0.0, 0.0, 1.0),
                    };

                    // Compute source contribution to each surface
                    let source_pot = self.compute_source_potential(dipole_pos, &dipole_ori);

                    // Solve BEM system
                    let surface_pot = &bem_inv * &source_pot;

                    // Interpolate to electrodes
                    for (s_idx, electrode_pos) in sensor_positions.iter().enumerate() {
                        let pot = self.interpolate_potential(&surface_pot, electrode_pos);
                        gain_data[s_idx][3 * src_idx + ori_idx] = pot;
                    }
                }
            } else {
                let dipole_ori = source
                    .orientation()
                    .unwrap_or_else(|| Vector3::new(0.0, 0.0, 1.0));

                let source_pot = self.compute_source_potential(dipole_pos, &dipole_ori);
                let surface_pot = &bem_inv * &source_pot;

                for (s_idx, electrode_pos) in sensor_positions.iter().enumerate() {
                    let pot = self.interpolate_potential(&surface_pot, electrode_pos);
                    gain_data[s_idx][src_idx] = pot;
                }
            }
        }

        let sensor_names: Vec<String> = sensors
            .iter()
            .filter(|s| matches!(s.sensor_type(), SensorType::Eeg))
            .map(|s| s.name().to_string())
            .collect();

        GainMatrix::new(gain_data, free_ori, sensor_names)
    }

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

    /// Get total number of triangles
    pub fn n_triangles(&self) -> usize {
        self.surfaces.iter().map(BemSurface::n_triangles).sum()
    }

    /// Get conductivities
    pub fn conductivities(&self) -> &[f64; 3] {
        &self.config.conductivities
    }

    // ========== Private methods ==========

    /// Compute BEM interaction block between two surfaces
    fn compute_bem_block(
        &self,
        surf_i: &BemSurface,
        surf_j: &BemSurface,
        same: bool,
    ) -> DMatrix<f64> {
        let ni = surf_i.n_triangles();
        let nj = surf_j.n_triangles();
        let mut block = DMatrix::zeros(ni, nj);

        for (i, ((ri, _ni), _ai)) in surf_i
            .centroids()
            .iter()
            .zip(surf_i.normals().iter())
            .zip(surf_i.areas().iter())
            .enumerate()
        {
            for (j, ((rj, nj_vec), aj)) in surf_j
                .centroids()
                .iter()
                .zip(surf_j.normals().iter())
                .zip(surf_j.areas().iter())
                .enumerate()
            {
                if same && i == j {
                    // Self-term: solid angle / (2π)
                    block[(i, j)] = *aj / (2.0 * PI);
                } else {
                    // Off-diagonal: compute solid angle
                    let solid_angle = self.compute_solid_angle(ri, rj, nj_vec, *aj);
                    block[(i, j)] = solid_angle / (4.0 * PI);
                }
            }
        }

        block
    }

    /// Compute solid angle subtended by a triangle at a point
    fn compute_solid_angle(
        &self,
        point: &Position,
        center: &Position,
        normal: &Orientation,
        area: f64,
    ) -> f64 {
        let r = center - point;
        let r_norm = norm(&r);

        if r_norm < 1e-10 {
            return 2.0 * PI; // Point is at the triangle
        }

        // Approximate solid angle for small triangle
        let r_hat = r / r_norm;
        let cos_theta = dot(&r_hat, normal);

        area * cos_theta / (r_norm * r_norm)
    }

    /// Compute source contribution to surface potentials
    fn compute_source_potential(
        &self,
        dipole_pos: &Position,
        dipole_ori: &Orientation,
    ) -> DVector<f64> {
        let mut pot = DVector::zeros(self.n_collocation);
        let sigma_brain = self.config.conductivities[0];

        let mut idx = 0;
        for surface in &self.surfaces {
            for (centroid, normal) in surface.centroids().iter().zip(surface.normals().iter()) {
                // Infinite medium potential at the surface
                let r = centroid - dipole_pos;
                let r_norm = norm(&r);

                if r_norm > 1e-10 {
                    // V = (1 / 4πσ) * (p · r) / r³
                    let p_dot_r = dot(dipole_ori, &r);
                    pot[idx] = p_dot_r / (4.0 * PI * sigma_brain * r_norm.powi(3));

                    // Add normal derivative contribution
                    let dn = dot(&r, normal) / r_norm;
                    pot[idx] += dn * p_dot_r / (4.0 * PI * sigma_brain * r_norm.powi(3));
                }

                idx += 1;
            }
        }

        pot
    }

    /// Interpolate potential at electrode position from surface potentials
    fn interpolate_potential(&self, surface_pot: &DVector<f64>, electrode_pos: &Position) -> f64 {
        // Find nearest scalp triangle (scalp is the last surface)
        let scalp = &self.surfaces[2];
        let offset: usize = self.surfaces[0].n_triangles() + self.surfaces[1].n_triangles();

        let mut min_dist = f64::INFINITY;
        let mut nearest_idx = 0;

        for (i, centroid) in scalp.centroids().iter().enumerate() {
            let dist = norm(&(electrode_pos - centroid));
            if dist < min_dist {
                min_dist = dist;
                nearest_idx = i;
            }
        }

        // Simple nearest-neighbor interpolation
        // (A more sophisticated implementation would use barycentric interpolation)
        surface_pot[offset + nearest_idx]
    }

    /// Invert a matrix using SVD
    fn invert_matrix(m: &DMatrix<f64>) -> ForwardResult<DMatrix<f64>> {
        let svd = m.clone().svd(true, true);

        let u = svd
            .u
            .ok_or_else(|| ForwardError::ComputationError("SVD failed: no U matrix".to_string()))?;

        let vt = svd
            .v_t
            .ok_or_else(|| ForwardError::ComputationError("SVD failed: no V matrix".to_string()))?;

        let s = svd.singular_values;
        let n = s.len();

        // Compute pseudoinverse
        let tol = 1e-10 * s[0];
        let s_inv = DMatrix::from_diagonal(&DVector::from_fn(n, |i, _| {
            if s[i] > tol { 1.0 / s[i] } else { 0.0 }
        }));

        Ok(vt.transpose() * &s_inv * u.transpose())
    }
}

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

    #[test]
    fn test_bem_surface_sphere() {
        let surface = BemSurface::sphere([0.0, 0.0, 0.04], 0.08, 2, "test").unwrap();

        assert!(surface.n_triangles() > 0);
        assert!(surface.n_vertices() > 0);
        assert_eq!(surface.name(), "test");
    }

    #[test]
    fn test_bem_surface_properties() {
        let surface = BemSurface::sphere([0.0, 0.0, 0.0], 0.1, 1, "sphere").unwrap();

        // Check that we have centroids, normals, and areas
        assert_eq!(surface.centroids().len(), surface.n_triangles());
        assert_eq!(surface.normals().len(), surface.n_triangles());
        assert_eq!(surface.areas().len(), surface.n_triangles());

        // All normals should be unit vectors
        for n in surface.normals() {
            let len = norm(n);
            assert!(
                (len - 1.0).abs() < 1e-6 || len < 1e-10,
                "Normal length: {}",
                len
            );
        }

        // All areas should be positive
        for &a in surface.areas() {
            assert!(a > 0.0);
        }
    }

    #[test]
    fn test_bem_config() {
        let config = BemConfig::default();
        assert!((config.conductivities[0] - 0.33).abs() < 1e-10);
        assert!((config.conductivities[1] - 0.0042).abs() < 1e-10);
        assert!((config.conductivities[2] - 0.33).abs() < 1e-10);
    }

    #[test]
    fn test_bem_model_creation() {
        let model = BemModel::three_shell_sphere(
            [0.0, 0.0, 0.04],
            [0.06, 0.065, 0.08], // inner_skull, outer_skull, scalp
            80,
            None,
        )
        .unwrap();

        assert_eq!(model.n_surfaces(), 3);
        assert!(model.n_triangles() > 0);
    }

    #[test]
    fn test_bem_matrix_computation() {
        let mut model =
            BemModel::three_shell_sphere([0.0, 0.0, 0.04], [0.06, 0.065, 0.08], 80, None).unwrap();

        // This should complete without error
        model.make_bem_matrix().unwrap();

        assert!(model.bem_matrix.is_some());
    }

    #[test]
    fn test_bem_gain_matrix() {
        let mut model =
            BemModel::three_shell_sphere([0.0, 0.0, 0.04], [0.06, 0.065, 0.08], 80, None).unwrap();

        model.make_bem_matrix().unwrap();

        // Create EEG electrode array
        let sensors = SensorArray::eeg_10_20(0.085);

        // Create simple source space
        let sources = SourceSpace::create_spherical_shell([0.0, 0.0, 0.04], 0.03, 10);

        let gain = model.compute_gain(&sources, &sensors).unwrap();

        assert_eq!(gain.n_sensors(), sensors.len());
        assert_eq!(gain.n_source_columns(), sources.len()); // Fixed orientation
    }
}