//! Lead field (gain matrix) computation for MEG/EEG forward modeling.
//!
//! The lead field matrix relates neural source currents to measured
//! sensor signals: M = L * J, where L is the lead field.

use crate::assembly::{FemAssembler, TransferMatrix};
use crate::error::{FemError, FemResult};
use crate::mesh::HeadMesh;
use crate::solver::{FemSolver, SolverConfig};
use nalgebra::Vector3;
use ndarray::{Array1, Array2, Axis};
use rayon::prelude::*;
use serde::{Deserialize, Serialize};
use std::time::Instant;

/// Source type for lead field computation
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum SourceType {
    /// Current dipole
    Dipole,
    /// Monopole (point source)
    Monopole,
}

/// Lead field configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LeadFieldConfig {
    /// Source type
    pub source_type: SourceType,
    /// Number of orientations per source (1 for fixed, 3 for free)
    pub n_orientations: usize,
    /// Whether to compute MEG lead field
    pub compute_meg: bool,
    /// Whether to compute EEG lead field
    pub compute_eeg: bool,
    /// Apply average reference to EEG
    pub eeg_average_reference: bool,
    /// Solver configuration
    pub solver_config: SolverConfig,
    /// Parallel computation
    pub parallel: bool,
    /// Verbose output
    pub verbose: bool,
}

impl Default for LeadFieldConfig {
    fn default() -> Self {
        Self {
            source_type: SourceType::Dipole,
            n_orientations: 3,
            compute_meg: false,
            compute_eeg: true,
            eeg_average_reference: true,
            solver_config: SolverConfig::default(),
            parallel: true,
            verbose: false,
        }
    }
}

/// Source space definition
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SourceSpace {
    /// Source positions (n_sources x 3)
    pub positions: Array2<f64>,
    /// Source normals/orientations (n_sources x 3, optional)
    pub normals: Option<Array2<f64>>,
    /// Element indices containing each source
    pub element_indices: Vec<Option<usize>>,
}

impl SourceSpace {
    /// Create source space from positions
    pub fn from_positions(positions: Array2<f64>) -> Self {
        let n_sources = positions.nrows();
        Self {
            positions,
            normals: None,
            element_indices: vec![None; n_sources],
        }
    }

    /// Create source space with normals
    pub fn with_normals(positions: Array2<f64>, normals: Array2<f64>) -> FemResult<Self> {
        if positions.nrows() != normals.nrows() {
            return Err(FemError::DimensionMismatch(
                "Positions and normals must have same number of rows".into(),
            ));
        }

        let n_sources = positions.nrows();
        Ok(Self {
            positions,
            normals: Some(normals),
            element_indices: vec![None; n_sources],
        })
    }

    /// Create regular grid source space
    pub fn regular_grid(center: &Vector3<f64>, extent: &Vector3<f64>, spacing: f64) -> Self {
        let nx = (extent.x / spacing).ceil() as usize;
        let ny = (extent.y / spacing).ceil() as usize;
        let nz = (extent.z / spacing).ceil() as usize;

        let n_sources = nx * ny * nz;
        let mut positions = Array2::zeros((n_sources, 3));

        let mut idx = 0;
        for i in 0..nx {
            for j in 0..ny {
                for k in 0..nz {
                    let x = center.x - extent.x / 2.0 + i as f64 * spacing;
                    let y = center.y - extent.y / 2.0 + j as f64 * spacing;
                    let z = center.z - extent.z / 2.0 + k as f64 * spacing;

                    positions[[idx, 0]] = x;
                    positions[[idx, 1]] = y;
                    positions[[idx, 2]] = z;
                    idx += 1;
                }
            }
        }

        Self::from_positions(positions)
    }

    /// Locate sources in mesh elements
    pub fn locate_in_mesh(&mut self, mesh: &HeadMesh) {
        for (i, row) in self.positions.axis_iter(Axis(0)).enumerate() {
            let pos = Vector3::new(row[0], row[1], row[2]);
            self.element_indices[i] = mesh.find_element(&pos);
        }
    }

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

    /// Get source position
    pub fn position(&self, idx: usize) -> Vector3<f64> {
        Vector3::new(
            self.positions[[idx, 0]],
            self.positions[[idx, 1]],
            self.positions[[idx, 2]],
        )
    }

    /// Get source normal (or default)
    pub fn normal(&self, idx: usize) -> Vector3<f64> {
        if let Some(ref normals) = self.normals {
            Vector3::new(normals[[idx, 0]], normals[[idx, 1]], normals[[idx, 2]])
        } else {
            Vector3::new(0.0, 0.0, 1.0)
        }
    }
}

/// EEG/MEG sensor positions
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SensorArray {
    /// Sensor positions (n_sensors x 3)
    pub positions: Array2<f64>,
    /// Sensor orientations for MEG (n_sensors x 3)
    pub orientations: Option<Array2<f64>>,
    /// Sensor labels
    pub labels: Vec<String>,
    /// Sensor type
    pub sensor_type: SensorType,
}

/// Sensor type
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum SensorType {
    /// EEG electrodes
    EEG,
    /// MEG magnetometers
    Magnetometer,
    /// MEG gradiometers
    Gradiometer,
}

impl SensorArray {
    /// Create EEG sensor array
    pub fn eeg(positions: Array2<f64>, labels: Vec<String>) -> Self {
        Self {
            positions,
            orientations: None,
            labels,
            sensor_type: SensorType::EEG,
        }
    }

    /// Create MEG sensor array
    pub fn meg(
        positions: Array2<f64>,
        orientations: Array2<f64>,
        labels: Vec<String>,
        sensor_type: SensorType,
    ) -> Self {
        Self {
            positions,
            orientations: Some(orientations),
            labels,
            sensor_type,
        }
    }

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

    /// Get sensor positions as vectors
    pub fn position_vectors(&self) -> Vec<Vector3<f64>> {
        (0..self.n_sensors())
            .map(|i| {
                Vector3::new(
                    self.positions[[i, 0]],
                    self.positions[[i, 1]],
                    self.positions[[i, 2]],
                )
            })
            .collect()
    }
}

/// Computed lead field matrix
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LeadField {
    /// Lead field matrix (n_sensors x n_sources * n_orientations)
    pub matrix: Array2<f64>,
    /// Number of sensors
    pub n_sensors: usize,
    /// Number of sources
    pub n_sources: usize,
    /// Number of orientations per source
    pub n_orientations: usize,
    /// Source positions
    pub source_positions: Array2<f64>,
    /// Sensor positions
    pub sensor_positions: Array2<f64>,
    /// Computation time (seconds)
    pub computation_time: f64,
    /// Sources inside mesh
    pub sources_in_mesh: usize,
}

impl LeadField {
    /// Get lead field for a single source and orientation
    pub fn get_column(&self, source_idx: usize, orientation: usize) -> Array1<f64> {
        let col_idx = source_idx * self.n_orientations + orientation;
        self.matrix.column(col_idx).to_owned()
    }

    /// Get lead field for all orientations of a source
    pub fn get_source(&self, source_idx: usize) -> Array2<f64> {
        let start = source_idx * self.n_orientations;
        let end = start + self.n_orientations;
        self.matrix.slice(ndarray::s![.., start..end]).to_owned()
    }

    /// Apply lead field: M = L * J
    pub fn apply(&self, source_currents: &Array1<f64>) -> FemResult<Array1<f64>> {
        let expected_len = self.n_sources * self.n_orientations;
        if source_currents.len() != expected_len {
            return Err(FemError::DimensionMismatch(format!(
                "Source currents length {} doesn't match expected {}",
                source_currents.len(),
                expected_len
            )));
        }

        Ok(self.matrix.dot(source_currents))
    }
}

/// Lead field computer
pub struct LeadFieldComputer<'a> {
    /// Configuration
    config: LeadFieldConfig,
    /// FEM assembler
    assembler: &'a FemAssembler,
    /// FEM solver
    solver: FemSolver,
    /// Transfer matrix for EEG
    eeg_transfer: Option<TransferMatrix>,
}

impl<'a> LeadFieldComputer<'a> {
    /// Create lead field computer
    pub fn new(assembler: &'a FemAssembler, config: LeadFieldConfig) -> FemResult<Self> {
        let solver = FemSolver::from_assembler(assembler, config.solver_config.clone())?;

        Ok(Self {
            config,
            assembler,
            solver,
            eeg_transfer: None,
        })
    }

    /// Set EEG electrode positions
    pub fn set_eeg_electrodes(&mut self, electrodes: &[Vector3<f64>]) -> FemResult<()> {
        self.eeg_transfer = Some(TransferMatrix::compute(self.assembler.mesh(), electrodes)?);
        Ok(())
    }

    /// Compute EEG lead field
    pub fn compute_eeg(
        &self,
        sources: &SourceSpace,
        sensors: &SensorArray,
    ) -> FemResult<LeadField> {
        let start_time = Instant::now();

        let n_sensors = sensors.n_sensors();
        let n_sources = sources.n_sources();
        let n_orient = self.config.n_orientations;
        let n_cols = n_sources * n_orient;

        if self.config.verbose {
            eprintln!(
                "Computing EEG lead field: {} sensors, {} sources, {} orientations",
                n_sensors, n_sources, n_orient
            );
        }

        // Compute transfer matrix if not already set
        let transfer = if let Some(ref t) = self.eeg_transfer {
            t.clone()
        } else {
            TransferMatrix::compute(self.assembler.mesh(), &sensors.position_vectors())?
        };

        // Compute lead field columns
        let mut leadfield = Array2::zeros((n_sensors, n_cols));
        let mut sources_in_mesh = 0;

        // Define orientation vectors
        let orientations: Vec<Vector3<f64>> = if n_orient == 1 {
            vec![Vector3::new(0.0, 0.0, 1.0)]
        } else {
            vec![
                Vector3::new(1.0, 0.0, 0.0),
                Vector3::new(0.0, 1.0, 0.0),
                Vector3::new(0.0, 0.0, 1.0),
            ]
        };

        if self.config.parallel {
            // Parallel computation over sources
            let results: Vec<(usize, Vec<Array1<f64>>)> = (0..n_sources)
                .into_par_iter()
                .filter_map(|src_idx| {
                    let pos = sources.position(src_idx);

                    // Find element containing source
                    let elem_idx = match sources.element_indices.get(src_idx) {
                        Some(Some(idx)) => *idx,
                        _ => self.assembler.mesh().find_element(&pos)?,
                    };

                    // Compute for each orientation
                    let mut cols = Vec::with_capacity(n_orient);

                    for orient in &orientations {
                        // Create dipole RHS
                        let rhs = self.create_dipole_rhs(src_idx, elem_idx, orient);

                        // Solve
                        if let Ok(result) = self.solver.solve(&rhs) {
                            let nodal_potential = Array1::from_vec(result.solution);
                            let sensor_potential = transfer.apply(&nodal_potential);
                            cols.push(sensor_potential);
                        } else {
                            cols.push(Array1::zeros(n_sensors));
                        }
                    }

                    Some((src_idx, cols))
                })
                .collect();

            // Assemble results
            for (src_idx, cols) in results {
                sources_in_mesh += 1;
                for (orient_idx, col) in cols.iter().enumerate() {
                    let col_idx = src_idx * n_orient + orient_idx;
                    for (row_idx, &val) in col.iter().enumerate() {
                        leadfield[[row_idx, col_idx]] = val;
                    }
                }
            }
        } else {
            // Sequential computation
            for src_idx in 0..n_sources {
                let pos = sources.position(src_idx);

                // Find element
                let elem_idx = match sources.element_indices.get(src_idx) {
                    Some(Some(idx)) => *idx,
                    _ => {
                        if let Some(idx) = self.assembler.mesh().find_element(&pos) {
                            idx
                        } else {
                            continue;
                        }
                    }
                };

                sources_in_mesh += 1;

                for (orient_idx, orient) in orientations.iter().enumerate() {
                    let rhs = self.create_dipole_rhs(src_idx, elem_idx, orient);

                    if let Ok(result) = self.solver.solve(&rhs) {
                        let nodal_potential = Array1::from_vec(result.solution);
                        let sensor_potential = transfer.apply(&nodal_potential);

                        let col_idx = src_idx * n_orient + orient_idx;
                        for (row_idx, &val) in sensor_potential.iter().enumerate() {
                            leadfield[[row_idx, col_idx]] = val;
                        }
                    }
                }

                if self.config.verbose && src_idx % 100 == 0 {
                    eprintln!("Computed source {}/{}", src_idx + 1, n_sources);
                }
            }
        }

        // Apply average reference
        if self.config.eeg_average_reference {
            leadfield = self.apply_average_reference(leadfield);
        }

        let computation_time = start_time.elapsed().as_secs_f64();

        if self.config.verbose {
            eprintln!(
                "Lead field computed in {:.2}s ({} sources in mesh)",
                computation_time, sources_in_mesh
            );
        }

        Ok(LeadField {
            matrix: leadfield,
            n_sensors,
            n_sources,
            n_orientations: n_orient,
            source_positions: sources.positions.clone(),
            sensor_positions: sensors.positions.clone(),
            computation_time,
            sources_in_mesh,
        })
    }

    /// Create RHS vector for a dipole at given position with given orientation
    fn create_dipole_rhs(
        &self,
        _source_idx: usize,
        element_idx: usize,
        orientation: &Vector3<f64>,
    ) -> Array1<f64> {
        let n_nodes = self.assembler.n_nodes();
        let mut rhs = Array1::zeros(n_nodes);

        let mesh = self.assembler.mesh();
        let elem = &mesh.elements[element_idx];
        let grads = elem.shape_gradients(&mesh.nodes);

        // Dipole source: ∫ J_p · ∇φ dV = J_p · ∇φ * V (constant for linear tet)
        // RHS contribution to node i: q * orientation · grad_i * volume
        let q = 1.0; // Unit dipole moment

        for (local_idx, &global_idx) in elem.nodes.iter().enumerate() {
            let grad = &grads[local_idx];
            rhs[global_idx] = q * orientation.dot(grad) * elem.volume;
        }

        rhs
    }

    /// Apply average reference to lead field
    fn apply_average_reference(&self, mut leadfield: Array2<f64>) -> Array2<f64> {
        let n_sensors = leadfield.nrows();

        if n_sensors == 0 {
            return leadfield;
        }

        // Subtract mean across sensors for each column
        for mut col in leadfield.axis_iter_mut(Axis(1)) {
            let mean = col.sum() / n_sensors as f64;
            col -= mean;
        }

        leadfield
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::assembly::FemAssembler;
    use crate::conductivity::TissueConductivity;
    use crate::mesh::HeadMesh;

    fn create_test_assembler() -> FemAssembler {
        let mesh = HeadMesh::three_layer_sphere(0.08, 0.007, 0.006, 2, 1).unwrap();
        let conductivity = TissueConductivity::default_isotropic();
        let mut assembler = FemAssembler::with_defaults(mesh, conductivity);
        assembler.assemble_global().unwrap();
        assembler
    }

    #[test]
    fn test_source_space_creation() {
        let mut positions = Array2::zeros((10, 3));
        for i in 0..10 {
            positions[[i, 0]] = i as f64 * 0.01;
            positions[[i, 1]] = 0.0;
            positions[[i, 2]] = 0.0;
        }

        let sources = SourceSpace::from_positions(positions);
        assert_eq!(sources.n_sources(), 10);
    }

    #[test]
    fn test_source_space_grid() {
        let center = Vector3::new(0.0, 0.0, 0.0);
        let extent = Vector3::new(0.04, 0.04, 0.04);
        let spacing = 0.01;

        let sources = SourceSpace::regular_grid(&center, &extent, spacing);
        assert!(sources.n_sources() > 0);
    }

    #[test]
    fn test_sensor_array() {
        let mut positions = Array2::zeros((19, 3));
        let labels: Vec<String> = (0..19).map(|i| format!("E{}", i)).collect();

        let sensors = SensorArray::eeg(positions, labels);
        assert_eq!(sensors.n_sensors(), 19);
        assert_eq!(sensors.sensor_type, SensorType::EEG);
    }

    #[test]
    fn test_leadfield_computer_creation() {
        let assembler = create_test_assembler();
        let config = LeadFieldConfig::default();

        let computer = LeadFieldComputer::new(&assembler, config);
        assert!(computer.is_ok());
    }

    #[test]
    fn test_create_dipole_rhs() {
        let assembler = create_test_assembler();
        let config = LeadFieldConfig::default();
        let computer = LeadFieldComputer::new(&assembler, config).unwrap();

        let orientation = Vector3::new(0.0, 0.0, 1.0);
        let rhs = computer.create_dipole_rhs(0, 0, &orientation);

        assert_eq!(rhs.len(), assembler.n_nodes());

        // RHS should be sparse (only 4 nodes per element)
        let nonzeros = rhs.iter().filter(|&&v| v.abs() > 1e-15).count();
        assert!(nonzeros <= 4);
    }

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

        let mut positions = Array2::zeros((5, 3));
        // Points inside mesh
        positions[[0, 2]] = 0.02; // Center-ish
        positions[[1, 0]] = 0.04;
        positions[[2, 1]] = 0.03;
        // Point outside
        positions[[3, 2]] = 0.5;
        // Another inside
        positions[[4, 2]] = 0.05;

        let mut sources = SourceSpace::from_positions(positions);
        sources.locate_in_mesh(&mesh);

        // Some should be found
        let found = sources
            .element_indices
            .iter()
            .filter(|x| x.is_some())
            .count();
        assert!(found >= 2);
    }

    #[test]
    fn test_average_reference() {
        let assembler = create_test_assembler();
        let config = LeadFieldConfig {
            eeg_average_reference: true,
            ..Default::default()
        };
        let computer = LeadFieldComputer::new(&assembler, config).unwrap();

        // Create test matrix
        let mut matrix = Array2::zeros((10, 5));
        for i in 0..10 {
            for j in 0..5 {
                matrix[[i, j]] = (i + j) as f64;
            }
        }

        let ref_matrix = computer.apply_average_reference(matrix);

        // Each column should have zero mean
        for j in 0..5 {
            let col_sum: f64 = ref_matrix.column(j).sum();
            assert!(col_sum.abs() < 1e-10, "Column {} sum = {}", j, col_sum);
        }
    }

    #[test]
    fn test_leadfield_apply() {
        // Create a simple lead field
        let matrix = Array2::from_shape_vec(
            (3, 6),
            vec![
                1.0, 0.0, 0.0, 2.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 2.0, 0.0, 0.0, 0.0, 1.0, 0.0,
                0.0, 2.0,
            ],
        )
        .unwrap();

        let leadfield = LeadField {
            matrix,
            n_sensors: 3,
            n_sources: 2,
            n_orientations: 3,
            source_positions: Array2::zeros((2, 3)),
            sensor_positions: Array2::zeros((3, 3)),
            computation_time: 0.0,
            sources_in_mesh: 2,
        };

        // Apply to source currents
        let currents = Array1::from_vec(vec![1.0, 0.0, 0.0, 0.0, 1.0, 0.0]);
        let sensors = leadfield.apply(&currents).unwrap();

        assert_eq!(sensors.len(), 3);
        assert!((sensors[0] - 1.0).abs() < 1e-10);
        assert!((sensors[1] - 2.0).abs() < 1e-10);
        assert!((sensors[2] - 0.0).abs() < 1e-10);
    }

    #[test]
    fn test_leadfield_get_source() {
        let matrix = Array2::from_shape_vec(
            (2, 6),
            vec![
                1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0,
            ],
        )
        .unwrap();

        let leadfield = LeadField {
            matrix,
            n_sensors: 2,
            n_sources: 2,
            n_orientations: 3,
            source_positions: Array2::zeros((2, 3)),
            sensor_positions: Array2::zeros((2, 3)),
            computation_time: 0.0,
            sources_in_mesh: 2,
        };

        // Get first source
        let src0 = leadfield.get_source(0);
        assert_eq!(src0.dim(), (2, 3));
        assert!((src0[[0, 0]] - 1.0).abs() < 1e-10);
        assert!((src0[[0, 2]] - 3.0).abs() < 1e-10);

        // Get second source
        let src1 = leadfield.get_source(1);
        assert!((src1[[0, 0]] - 4.0).abs() < 1e-10);
    }
}