//! Head geometry and sensor array models for MEG/EEG.
//!
//! Provides head models and sensor configurations for PINN-based
//! source localization.

use crate::conductivity::{ConductivityModel, LayeredConductivity, TissueLayer, TissueType};
use crate::error::{PinnError, PinnResult};
use ndarray::{Array1, Array2};
use serde::{Deserialize, Serialize};

/// Head geometry representation
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum HeadGeometry {
    /// Spherical head model
    Spherical {
        /// Center of sphere
        center: [f64; 3],
        /// Outer radius (scalp surface)
        radius: f64,
    },
    /// Realistic mesh-based model
    Mesh {
        /// Vertices [n_vertices, 3]
        vertices: Array2<f64>,
        /// Triangular faces [n_faces, 3]
        faces: Array2<u32>,
    },
}

impl HeadGeometry {
    /// Create a spherical head geometry
    pub fn sphere(center: [f64; 3], radius: f64) -> Self {
        HeadGeometry::Spherical { center, radius }
    }

    /// Check if a point is inside the head
    pub fn contains(&self, point: &[f64; 3]) -> bool {
        match self {
            HeadGeometry::Spherical { center, radius } => {
                let dx = point[0] - center[0];
                let dy = point[1] - center[1];
                let dz = point[2] - center[2];
                let r = (dx * dx + dy * dy + dz * dz).sqrt();
                r <= *radius
            }
            HeadGeometry::Mesh { vertices, .. } => {
                // Simple bounding box check for now
                let min_x = vertices.column(0).fold(f64::MAX, |a, &b| a.min(b));
                let max_x = vertices.column(0).fold(f64::MIN, |a, &b| a.max(b));
                let min_y = vertices.column(1).fold(f64::MAX, |a, &b| a.min(b));
                let max_y = vertices.column(1).fold(f64::MIN, |a, &b| a.max(b));
                let min_z = vertices.column(2).fold(f64::MAX, |a, &b| a.min(b));
                let max_z = vertices.column(2).fold(f64::MIN, |a, &b| a.max(b));

                point[0] >= min_x
                    && point[0] <= max_x
                    && point[1] >= min_y
                    && point[1] <= max_y
                    && point[2] >= min_z
                    && point[2] <= max_z
            }
        }
    }

    /// Get center of head
    pub fn center(&self) -> [f64; 3] {
        match self {
            HeadGeometry::Spherical { center, .. } => *center,
            HeadGeometry::Mesh { vertices, .. } => {
                let cx = vertices.column(0).mean().unwrap_or(0.0);
                let cy = vertices.column(1).mean().unwrap_or(0.0);
                let cz = vertices.column(2).mean().unwrap_or(0.0);
                [cx, cy, cz]
            }
        }
    }

    /// Get approximate radius
    pub fn approximate_radius(&self) -> f64 {
        match self {
            HeadGeometry::Spherical { radius, .. } => *radius,
            HeadGeometry::Mesh { vertices, .. } => {
                let center = self.center();
                vertices
                    .rows()
                    .into_iter()
                    .map(|row| {
                        let dx = row[0] - center[0];
                        let dy = row[1] - center[1];
                        let dz = row[2] - center[2];
                        (dx * dx + dy * dy + dz * dz).sqrt()
                    })
                    .fold(0.0f64, f64::max)
            }
        }
    }
}

/// Sensor modality
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum SensorType {
    /// EEG electrode
    Eeg,
    /// MEG gradiometer
    MegGrad,
    /// MEG magnetometer
    MegMag,
}

/// A single sensor
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Sensor {
    /// Sensor name
    pub name: String,
    /// Sensor type
    pub sensor_type: SensorType,
    /// Position [x, y, z] in meters
    pub position: [f64; 3],
    /// Orientation (for MEG) - unit vector
    pub orientation: Option<[f64; 3]>,
}

impl Sensor {
    /// Create an EEG electrode
    pub fn eeg(name: impl Into<String>, position: [f64; 3]) -> Self {
        Self {
            name: name.into(),
            sensor_type: SensorType::Eeg,
            position,
            orientation: None,
        }
    }

    /// Create a MEG magnetometer
    pub fn meg_mag(name: impl Into<String>, position: [f64; 3], orientation: [f64; 3]) -> Self {
        // Normalize orientation
        let len = (orientation[0].powi(2) + orientation[1].powi(2) + orientation[2].powi(2)).sqrt();
        let norm = if len > 1e-10 {
            [
                orientation[0] / len,
                orientation[1] / len,
                orientation[2] / len,
            ]
        } else {
            [0.0, 0.0, 1.0]
        };

        Self {
            name: name.into(),
            sensor_type: SensorType::MegMag,
            position,
            orientation: Some(norm),
        }
    }

    /// Create a MEG gradiometer
    pub fn meg_grad(name: impl Into<String>, position: [f64; 3], orientation: [f64; 3]) -> Self {
        let len = (orientation[0].powi(2) + orientation[1].powi(2) + orientation[2].powi(2)).sqrt();
        let norm = if len > 1e-10 {
            [
                orientation[0] / len,
                orientation[1] / len,
                orientation[2] / len,
            ]
        } else {
            [0.0, 0.0, 1.0]
        };

        Self {
            name: name.into(),
            sensor_type: SensorType::MegGrad,
            position,
            orientation: Some(norm),
        }
    }
}

/// Sensor array configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SensorArray {
    /// List of sensors
    sensors: Vec<Sensor>,
    /// Sensor positions as array [n_sensors, 3]
    positions: Array2<f64>,
    /// Sensor orientations as array [n_sensors, 3] (for MEG)
    orientations: Option<Array2<f64>>,
}

impl SensorArray {
    /// Create from list of sensors
    pub fn new(sensors: Vec<Sensor>) -> Self {
        let n = sensors.len();
        let mut positions = Array2::zeros((n, 3));
        let mut orientations = Array2::zeros((n, 3));
        let mut has_orientations = false;

        for (i, sensor) in sensors.iter().enumerate() {
            positions[[i, 0]] = sensor.position[0];
            positions[[i, 1]] = sensor.position[1];
            positions[[i, 2]] = sensor.position[2];

            if let Some(orient) = sensor.orientation {
                orientations[[i, 0]] = orient[0];
                orientations[[i, 1]] = orient[1];
                orientations[[i, 2]] = orient[2];
                has_orientations = true;
            }
        }

        Self {
            sensors,
            positions,
            orientations: if has_orientations {
                Some(orientations)
            } else {
                None
            },
        }
    }

    /// Create standard 10-20 EEG montage on spherical head
    pub fn standard_1020(head_radius: f64) -> Self {
        let electrode_names = [
            "Fp1", "Fp2", "F7", "F3", "Fz", "F4", "F8", "T3", "C3", "Cz", "C4", "T4", "T5", "P3",
            "Pz", "P4", "T6", "O1", "Oz", "O2",
        ];

        // Approximate 10-20 positions (normalized to unit sphere)
        let theta_phi = [
            (0.3, -0.4),
            (0.3, 0.4), // Fp1, Fp2
            (0.6, -0.8),
            (0.5, -0.4),
            (0.5, 0.0),
            (0.5, 0.4),
            (0.6, 0.8), // F7-F8
            (1.0, -1.0),
            (0.9, -0.4),
            (0.9, 0.0),
            (0.9, 0.4),
            (1.0, 1.0), // T3-T4
            (1.4, -0.8),
            (1.3, -0.4),
            (1.3, 0.0),
            (1.3, 0.4),
            (1.4, 0.8), // T5-T6
            (1.7, -0.4),
            (1.7, 0.0),
            (1.7, 0.4), // O1, Oz, O2
        ];

        let sensors: Vec<Sensor> = electrode_names
            .iter()
            .zip(theta_phi.iter())
            .map(|(&name, &(theta, phi)): (&&str, &(f64, f64))| {
                let x = head_radius * theta.sin() * phi.cos();
                let y = head_radius * theta.sin() * phi.sin();
                let z = head_radius * theta.cos();
                Sensor::eeg(name, [x, y, z])
            })
            .collect();

        SensorArray::new(sensors)
    }

    /// Create MEG helmet array
    pub fn meg_helmet(n_sensors: usize, head_radius: f64, helmet_distance: f64) -> Self {
        let sensor_radius = head_radius + helmet_distance;
        let mut sensors = Vec::with_capacity(n_sensors);

        // Distribute sensors using Fibonacci sphere
        let golden_ratio = f64::midpoint(1.0, 5.0_f64.sqrt());

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

            // Only upper hemisphere (z > 0)
            if phi < std::f64::consts::PI / 2.0 + 0.2 {
                let x = sensor_radius * phi.sin() * theta.cos();
                let y = sensor_radius * phi.sin() * theta.sin();
                let z = sensor_radius * phi.cos();

                // Orientation pointing toward head center
                let len = (x * x + y * y + z * z).sqrt();
                let orient = [-x / len, -y / len, -z / len];

                sensors.push(Sensor::meg_mag(format!("MEG{:03}", i), [x, y, z], orient));
            }
        }

        SensorArray::new(sensors)
    }

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

    /// Get sensor positions
    pub fn positions(&self) -> &Array2<f64> {
        &self.positions
    }

    /// Get sensor orientations
    pub fn orientations(&self) -> Option<&Array2<f64>> {
        self.orientations.as_ref()
    }

    /// Get sensor by name
    pub fn get(&self, name: &str) -> Option<&Sensor> {
        self.sensors.iter().find(|s| s.name == name)
    }

    /// Get sensor by index
    pub fn get_by_index(&self, idx: usize) -> Option<&Sensor> {
        self.sensors.get(idx)
    }

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

    /// Get sensor names
    pub fn names(&self) -> Vec<&str> {
        self.sensors.iter().map(|s| s.name.as_str()).collect()
    }
}

/// Complete head model with geometry, conductivity, and sensors
#[derive(Debug, Clone)]
pub struct HeadModel {
    /// Head geometry
    geometry: HeadGeometry,
    /// Tissue conductivity model
    conductivity: LayeredConductivity,
    /// Number of layers
    n_layers: usize,
}

impl HeadModel {
    /// Create a spherical head model
    pub fn spherical(n_layers: usize) -> Self {
        let center = [0.0, 0.0, 0.0];
        let geometry = HeadGeometry::sphere(center, 0.1); // 10 cm radius

        let conductivity = match n_layers {
            1 => {
                let mut model = LayeredConductivity::new(center);
                model.add_layer(
                    TissueLayer::new(TissueType::GrayMatter, "brain").with_radii(0.0, 0.1),
                );
                model
            }
            3 => LayeredConductivity::three_layer(center, 0.08, 0.007, 0.006),
            4 => LayeredConductivity::four_layer(center, 0.078, 0.002, 0.007, 0.006),
            _ => LayeredConductivity::three_layer(center, 0.08, 0.007, 0.006),
        };

        Self {
            geometry,
            conductivity,
            n_layers: n_layers.max(1).min(4),
        }
    }

    /// Set conductivity for a tissue type
    pub fn with_conductivity(mut self, tissue_name: &str, sigma: f64) -> Self {
        for layer in self.conductivity.layers_mut() {
            if layer.name == tissue_name {
                layer.conductivity = sigma;
            }
        }
        self
    }

    /// Make a tissue conductivity learnable
    pub fn with_learnable(mut self, tissue_name: &str) -> Self {
        for layer in self.conductivity.layers_mut() {
            if layer.name == tissue_name {
                layer.learnable = true;
            }
        }
        self
    }

    /// Get geometry
    pub fn geometry(&self) -> &HeadGeometry {
        &self.geometry
    }

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

    /// Get mutable conductivity model
    pub fn conductivity_mut(&mut self) -> &mut LayeredConductivity {
        &mut self.conductivity
    }

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

    /// Check if point is inside brain
    pub fn is_in_brain(&self, point: &[f64; 3]) -> bool {
        let layer = self.conductivity.layer_at(point);
        matches!(layer, Some(l) if l.tissue_type == TissueType::GrayMatter || l.tissue_type == TissueType::WhiteMatter)
    }

    /// Sample random points inside brain
    pub fn sample_brain_points(&self, n_points: usize) -> Array2<f64> {
        let mut points = Array2::zeros((n_points, 3));
        let center = self.geometry.center();

        // Get brain outer radius (first layer for spherical model)
        let brain_radius = self
            .conductivity
            .layers()
            .iter()
            .find(|l| l.tissue_type == TissueType::GrayMatter)
            .and_then(|l| l.outer_radius)
            .unwrap_or(0.08);

        // Simple uniform sampling in sphere
        let mut count = 0;
        let mut seed = 12345u64;

        while count < n_points {
            // LCG random number generator
            seed = seed.wrapping_mul(6364136223846793005).wrapping_add(1);
            let u1 = (seed as f64) / (u64::MAX as f64);
            seed = seed.wrapping_mul(6364136223846793005).wrapping_add(1);
            let u2 = (seed as f64) / (u64::MAX as f64);
            seed = seed.wrapping_mul(6364136223846793005).wrapping_add(1);
            let u3 = (seed as f64) / (u64::MAX as f64);

            // Map to [-1, 1]
            let x = (u1 * 2.0 - 1.0) * brain_radius;
            let y = (u2 * 2.0 - 1.0) * brain_radius;
            let z = (u3 * 2.0 - 1.0) * brain_radius;

            let r = (x * x + y * y + z * z).sqrt();
            if r <= brain_radius {
                points[[count, 0]] = center[0] + x;
                points[[count, 1]] = center[1] + y;
                points[[count, 2]] = center[2] + z;
                count += 1;
            }
        }

        points
    }

    /// Get conductivity at points
    pub fn conductivity_at_points(&self, points: &Array2<f64>) -> Array1<f64> {
        self.conductivity.conductivity_field(points)
    }
}

/// Source space (grid of potential source locations)
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SourceSpace {
    /// Source positions [n_sources, 3]
    pub positions: Array2<f64>,
    /// Source orientations [n_sources, 3] (optional, for fixed orientation)
    pub orientations: Option<Array2<f64>>,
    /// Whether orientations are fixed or free
    pub fixed_orientation: bool,
}

impl SourceSpace {
    /// Create source space from positions
    pub fn new(positions: Array2<f64>) -> Self {
        Self {
            positions,
            orientations: None,
            fixed_orientation: false,
        }
    }

    /// Create with fixed orientations
    pub fn with_orientations(
        positions: Array2<f64>,
        orientations: Array2<f64>,
    ) -> PinnResult<Self> {
        if positions.nrows() != orientations.nrows() {
            return Err(PinnError::DimensionMismatch(
                "Positions and orientations must have same number of sources".into(),
            ));
        }
        Ok(Self {
            positions,
            orientations: Some(orientations),
            fixed_orientation: true,
        })
    }

    /// Create grid source space inside brain
    pub fn grid(head: &HeadModel, spacing: f64) -> Self {
        let center = head.geometry().center();
        let radius = head.geometry().approximate_radius();

        // Find brain radius
        let brain_radius = head
            .conductivity()
            .layers()
            .iter()
            .find(|l| l.tissue_type == TissueType::GrayMatter)
            .and_then(|l| l.outer_radius)
            .unwrap_or(radius * 0.8);

        let n_per_dim = ((2.0 * brain_radius / spacing) as usize).max(1);
        let mut positions = Vec::new();

        for i in 0..n_per_dim {
            for j in 0..n_per_dim {
                for k in 0..n_per_dim {
                    let x = center[0] + (i as f64 - n_per_dim as f64 / 2.0) * spacing;
                    let y = center[1] + (j as f64 - n_per_dim as f64 / 2.0) * spacing;
                    let z = center[2] + (k as f64 - n_per_dim as f64 / 2.0) * spacing;

                    let r = ((x - center[0]).powi(2)
                        + (y - center[1]).powi(2)
                        + (z - center[2]).powi(2))
                    .sqrt();

                    if r <= brain_radius {
                        positions.push([x, y, z]);
                    }
                }
            }
        }

        let n = positions.len();
        let mut pos_array = Array2::zeros((n, 3));
        for (i, pos) in positions.iter().enumerate() {
            pos_array[[i, 0]] = pos[0];
            pos_array[[i, 1]] = pos[1];
            pos_array[[i, 2]] = pos[2];
        }

        SourceSpace::new(pos_array)
    }

    /// Number of sources
    pub fn n_sources(&self) -> usize {
        self.positions.nrows()
    }
}

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

    #[test]
    fn test_spherical_head() {
        let head = HeadModel::spherical(3);
        assert_eq!(head.n_layers(), 3);

        // Point inside brain
        assert!(head.is_in_brain(&[0.0, 0.0, 0.05]));

        // Point in skull
        assert!(!head.is_in_brain(&[0.0, 0.0, 0.085]));
    }

    #[test]
    fn test_sensor_array_1020() {
        let sensors = SensorArray::standard_1020(0.1);
        assert_eq!(sensors.n_sensors(), 20);
        assert!(sensors.get("Cz").is_some());
    }

    #[test]
    fn test_meg_helmet() {
        let sensors = SensorArray::meg_helmet(100, 0.1, 0.02);
        assert!(sensors.n_sensors() > 0);
        assert!(sensors.orientations().is_some());
    }

    #[test]
    fn test_sample_brain_points() {
        let head = HeadModel::spherical(3);
        let points = head.sample_brain_points(100);

        assert_eq!(points.nrows(), 100);
        assert_eq!(points.ncols(), 3);

        // All points should be inside brain
        for i in 0..100 {
            let point = [points[[i, 0]], points[[i, 1]], points[[i, 2]]];
            assert!(head.is_in_brain(&point));
        }
    }

    #[test]
    fn test_source_space_grid() {
        let head = HeadModel::spherical(3);
        let sources = SourceSpace::grid(&head, 0.02);

        assert!(sources.n_sources() > 0);
        assert_eq!(sources.positions.ncols(), 3);
    }

    #[test]
    fn test_geometry_contains() {
        let geom = HeadGeometry::sphere([0.0, 0.0, 0.0], 0.1);
        assert!(geom.contains(&[0.0, 0.0, 0.0]));
        assert!(geom.contains(&[0.05, 0.05, 0.05]));
        assert!(!geom.contains(&[0.2, 0.0, 0.0]));
    }
}