//! Maxwell equation residuals for bioelectromagnetic source localization.
//!
//! Implements the quasi-static Maxwell equations that govern neural current flow:
//!
//! ```text
//! ∇·(σ∇Φ) = ∇·Jp
//! ```
//!
//! where Φ is the electric potential, σ is the tissue conductivity,
//! and Jp is the primary (neural) current density.

use crate::error::{PinnError, PinnResult};
use ndarray::{Array1, Array2, Array3};
use serde::{Deserialize, Serialize};

/// Current density representation
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CurrentDensity {
    /// Current density components [n_sources, 3] (Jx, Jy, Jz)
    pub components: Array2<f64>,
    /// Source positions [n_sources, 3]
    pub positions: Array2<f64>,
    /// Magnitude at each source
    pub magnitudes: Array1<f64>,
}

impl CurrentDensity {
    /// Create a new current density field
    pub fn new(n_sources: usize) -> Self {
        Self {
            components: Array2::zeros((n_sources, 3)),
            positions: Array2::zeros((n_sources, 3)),
            magnitudes: Array1::zeros(n_sources),
        }
    }

    /// Create from positions and orientations
    pub fn from_dipoles(
        positions: Array2<f64>,
        orientations: Array2<f64>,
        magnitudes: Array1<f64>,
    ) -> PinnResult<Self> {
        let n = positions.nrows();
        if orientations.nrows() != n || magnitudes.len() != n {
            return Err(PinnError::DimensionMismatch(
                "Positions, orientations, and magnitudes must have same length".into(),
            ));
        }

        // Scale orientations by magnitudes
        let mut components = Array2::zeros((n, 3));
        for i in 0..n {
            for j in 0..3 {
                components[[i, j]] = orientations[[i, j]] * magnitudes[i];
            }
        }

        Ok(Self {
            components,
            positions,
            magnitudes,
        })
    }

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

    /// Compute divergence at a point (for physics residual)
    pub fn divergence_at(&self, point: &[f64; 3], sigma: f64) -> f64 {
        // Simplified: sum of contributions from all sources
        // In reality, this would involve Green's functions
        let mut div = 0.0;
        let eps = 1e-10;

        for i in 0..self.n_sources() {
            let dx = point[0] - self.positions[[i, 0]];
            let dy = point[1] - self.positions[[i, 1]];
            let dz = point[2] - self.positions[[i, 2]];
            let r2 = dx * dx + dy * dy + dz * dz + eps;
            let r = r2.sqrt();

            // Dipole-like contribution
            let jx = self.components[[i, 0]];
            let jy = self.components[[i, 1]];
            let jz = self.components[[i, 2]];

            // ∇·J contribution (simplified monopole approximation)
            let dot = jx * dx + jy * dy + jz * dz;
            div += dot / (r2 * r) / sigma;
        }

        div
    }

    /// Total current magnitude
    pub fn total_magnitude(&self) -> f64 {
        self.magnitudes.iter().map(|&m| m.abs()).sum()
    }
}

/// Quasi-static Maxwell equation for bioelectric problems
#[derive(Debug, Clone)]
pub struct QuasiStaticMaxwell {
    /// Default conductivity (S/m)
    conductivity: f64,
    /// Regularization parameter for sparsity
    sparsity_weight: f64,
    /// Regularization for smoothness
    smoothness_weight: f64,
}

impl QuasiStaticMaxwell {
    /// Create a new quasi-static Maxwell equation
    pub fn new(conductivity: f64) -> Self {
        Self {
            conductivity,
            sparsity_weight: 0.01,
            smoothness_weight: 0.001,
        }
    }

    /// Set sparsity regularization weight
    pub fn with_sparsity(mut self, weight: f64) -> Self {
        self.sparsity_weight = weight;
        self
    }

    /// Set smoothness regularization weight
    pub fn with_smoothness(mut self, weight: f64) -> Self {
        self.smoothness_weight = weight;
        self
    }

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

    /// Compute the PDE residual at collocation points
    ///
    /// The residual is: ∇·(σ∇Φ) - ∇·Jp = 0
    ///
    /// # Arguments
    /// * `potential` - Electric potential field [n_points]
    /// * `potential_grad` - Gradient of potential [n_points, 3]
    /// * `current_density` - Current density field
    /// * `conductivity_field` - Spatially varying conductivity [n_points]
    ///
    /// # Returns
    /// Residual at each collocation point
    pub fn compute_residual(
        &self,
        potential: &Array1<f64>,
        potential_grad: &Array2<f64>,
        potential_hessian: &Array3<f64>,
        current_density: &CurrentDensity,
        conductivity_field: &Array1<f64>,
        points: &Array2<f64>,
    ) -> PinnResult<Array1<f64>> {
        let n_points = potential.len();
        if potential_grad.nrows() != n_points {
            return Err(PinnError::DimensionMismatch(
                "Potential and gradient dimensions don't match".into(),
            ));
        }

        let mut residual = Array1::zeros(n_points);

        for i in 0..n_points {
            let sigma = conductivity_field[i];

            // Laplacian term: ∇²Φ (trace of Hessian)
            let laplacian = potential_hessian[[i, 0, 0]]
                + potential_hessian[[i, 1, 1]]
                + potential_hessian[[i, 2, 2]];

            // σ∇²Φ term
            let diffusion = sigma * laplacian;

            // ∇σ·∇Φ term (if conductivity varies spatially)
            // For now, assume locally constant conductivity
            let advection = 0.0;

            // Source term: ∇·Jp at this point
            let point = [points[[i, 0]], points[[i, 1]], points[[i, 2]]];
            let source = current_density.divergence_at(&point, sigma);

            // Residual: ∇·(σ∇Φ) - ∇·Jp should be zero
            residual[i] = diffusion + advection - source;
        }

        Ok(residual)
    }

    /// Compute physics loss from residuals
    pub fn physics_loss(&self, residual: &Array1<f64>) -> f64 {
        // Mean squared residual

        residual.iter().map(|&r| r * r).sum::<f64>() / residual.len() as f64
    }

    /// Compute sparsity loss for source estimation
    pub fn sparsity_loss(&self, current_density: &CurrentDensity) -> f64 {
        // L1 norm of source magnitudes
        let l1 = current_density
            .magnitudes
            .iter()
            .map(|&m| m.abs())
            .sum::<f64>();
        self.sparsity_weight * l1
    }

    /// Compute smoothness loss for source distribution
    pub fn smoothness_loss(&self, current_density: &CurrentDensity) -> f64 {
        if current_density.n_sources() < 2 {
            return 0.0;
        }

        // Total variation of source magnitudes (simplified)
        let mut tv = 0.0;
        for i in 1..current_density.n_sources() {
            let diff = current_density.magnitudes[i] - current_density.magnitudes[i - 1];
            tv += diff.abs();
        }

        self.smoothness_weight * tv
    }

    /// Compute total loss
    pub fn total_loss(&self, residual: &Array1<f64>, current_density: &CurrentDensity) -> f64 {
        self.physics_loss(residual)
            + self.sparsity_loss(current_density)
            + self.smoothness_loss(current_density)
    }
}

/// Maxwell residual trait for different formulations
pub trait MaxwellResidual {
    /// Compute residual at points
    fn residual(
        &self,
        potential: &Array1<f64>,
        gradients: &Array2<f64>,
        hessians: &Array3<f64>,
        sources: &CurrentDensity,
        conductivity: &Array1<f64>,
        points: &Array2<f64>,
    ) -> PinnResult<Array1<f64>>;

    /// Compute loss from residual
    fn loss(&self, residual: &Array1<f64>) -> f64;
}

impl MaxwellResidual for QuasiStaticMaxwell {
    fn residual(
        &self,
        potential: &Array1<f64>,
        gradients: &Array2<f64>,
        hessians: &Array3<f64>,
        sources: &CurrentDensity,
        conductivity: &Array1<f64>,
        points: &Array2<f64>,
    ) -> PinnResult<Array1<f64>> {
        self.compute_residual(
            potential,
            gradients,
            hessians,
            sources,
            conductivity,
            points,
        )
    }

    fn loss(&self, residual: &Array1<f64>) -> f64 {
        self.physics_loss(residual)
    }
}

/// MEG-specific Maxwell formulation (magnetic field)
#[derive(Debug, Clone)]
pub struct MagneticMaxwell {
    /// Permeability of free space
    mu0: f64,
}

impl MagneticMaxwell {
    /// Create new magnetic Maxwell formulation
    pub fn new() -> Self {
        Self {
            mu0: 4.0 * std::f64::consts::PI * 1e-7, // H/m
        }
    }

    /// Compute magnetic field from current density using Biot-Savart
    pub fn magnetic_field(
        &self,
        current_density: &CurrentDensity,
        sensor_positions: &Array2<f64>,
    ) -> Array2<f64> {
        let n_sensors = sensor_positions.nrows();
        let mut b_field = Array2::zeros((n_sensors, 3));

        for s in 0..n_sensors {
            let sensor = [
                sensor_positions[[s, 0]],
                sensor_positions[[s, 1]],
                sensor_positions[[s, 2]],
            ];

            for i in 0..current_density.n_sources() {
                let source = [
                    current_density.positions[[i, 0]],
                    current_density.positions[[i, 1]],
                    current_density.positions[[i, 2]],
                ];

                let j = [
                    current_density.components[[i, 0]],
                    current_density.components[[i, 1]],
                    current_density.components[[i, 2]],
                ];

                // r = sensor - source
                let r = [
                    sensor[0] - source[0],
                    sensor[1] - source[1],
                    sensor[2] - source[2],
                ];
                let r_mag = (r[0] * r[0] + r[1] * r[1] + r[2] * r[2]).sqrt();
                let r_mag3 = r_mag * r_mag * r_mag;

                if r_mag > 1e-10 {
                    // B = (μ0/4π) * (J × r) / |r|³
                    let cross = [
                        j[1] * r[2] - j[2] * r[1],
                        j[2] * r[0] - j[0] * r[2],
                        j[0] * r[1] - j[1] * r[0],
                    ];

                    let coeff = self.mu0 / (4.0 * std::f64::consts::PI * r_mag3);
                    b_field[[s, 0]] += coeff * cross[0];
                    b_field[[s, 1]] += coeff * cross[1];
                    b_field[[s, 2]] += coeff * cross[2];
                }
            }
        }

        b_field
    }

    /// Compute data fitting loss for MEG
    pub fn data_loss(&self, predicted: &Array2<f64>, measured: &Array2<f64>) -> PinnResult<f64> {
        if predicted.shape() != measured.shape() {
            return Err(PinnError::DimensionMismatch(
                "Predicted and measured have different shapes".into(),
            ));
        }

        let mse = predicted
            .iter()
            .zip(measured.iter())
            .map(|(&p, &m)| (p - m).powi(2))
            .sum::<f64>()
            / predicted.len() as f64;

        Ok(mse)
    }
}

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

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

    #[test]
    fn test_current_density_creation() {
        let positions =
            Array2::from_shape_vec((2, 3), vec![0.0, 0.0, 0.05, 0.01, 0.0, 0.05]).unwrap();
        let orientations =
            Array2::from_shape_vec((2, 3), vec![0.0, 0.0, 1.0, 0.0, 0.0, 1.0]).unwrap();
        let magnitudes = Array1::from_vec(vec![1e-9, 2e-9]);

        let current = CurrentDensity::from_dipoles(positions, orientations, magnitudes).unwrap();
        assert_eq!(current.n_sources(), 2);
    }

    #[test]
    fn test_quasi_static_maxwell() {
        let maxwell = QuasiStaticMaxwell::new(0.33);
        assert!((maxwell.conductivity() - 0.33).abs() < 1e-10);
    }

    #[test]
    fn test_magnetic_field() {
        let mag = MagneticMaxwell::new();

        // Single dipole at origin pointing in z
        let positions = Array2::from_shape_vec((1, 3), vec![0.0, 0.0, 0.0]).unwrap();
        let orientations = Array2::from_shape_vec((1, 3), vec![0.0, 0.0, 1.0]).unwrap();
        let magnitudes = Array1::from_vec(vec![1e-9]);
        let current = CurrentDensity::from_dipoles(positions, orientations, magnitudes).unwrap();

        // Sensor at (0.1, 0, 0)
        let sensors = Array2::from_shape_vec((1, 3), vec![0.1, 0.0, 0.0]).unwrap();
        let b_field = mag.magnetic_field(&current, &sensors);

        // B should be in y direction (perpendicular to both z and x)
        assert!(b_field[[0, 0]].abs() < 1e-20); // Bx ~ 0
        assert!(b_field[[0, 1]].abs() > 0.0); // By != 0
        assert!(b_field[[0, 2]].abs() < 1e-20); // Bz ~ 0
    }
}