//! Spherical MEG forward model using the Sarvas formula.
//!
//! The Sarvas formula provides an analytical solution for the magnetic field
//! produced by a current dipole in a spherically symmetric conductor.
//!
//! Reference: Sarvas, J. (1987). Basic mathematical and electromagnetic
//! concepts of the biomagnetic inverse problem. Physics in Medicine & Biology.

use crate::sensors::{MegCoil, Sensor, SensorArray};
use crate::source_space::SourceSpace;
use crate::{ForwardError, ForwardResult, Position, cross, dot, norm};
use nalgebra::Vector3;
use std::f64::consts::PI;

/// Magnetic permeability of free space (H/m)
const MU_0: f64 = 4.0 * PI * 1e-7;

/// Spherical MEG forward model
#[derive(Debug, Clone)]
pub struct SphericalMeg {
    /// Center of the sphere in meters
    origin: Position,
    /// Radius of the conducting sphere in meters
    radius: f64,
}

impl SphericalMeg {
    /// Create a new spherical MEG model
    ///
    /// # Arguments
    /// * `origin` - Center of the sphere [x, y, z] in meters
    /// * `radius` - Radius of the sphere in meters
    pub fn new(origin: [f64; 3], radius: f64) -> ForwardResult<Self> {
        if radius <= 0.0 {
            return Err(ForwardError::InvalidGeometry(
                "Sphere radius must be positive".to_string(),
            ));
        }

        Ok(Self {
            origin: Vector3::new(origin[0], origin[1], origin[2]),
            radius,
        })
    }

    /// Create a default head model (typical adult head)
    ///
    /// Origin at [0, 0, 0.04] m (4 cm above origin), radius 0.08 m (8 cm)
    pub fn default_head() -> Self {
        Self {
            origin: Vector3::new(0.0, 0.0, 0.04),
            radius: 0.08,
        }
    }

    /// Get the sphere origin
    pub fn origin(&self) -> &Position {
        &self.origin
    }

    /// Get the sphere radius
    pub fn radius(&self) -> f64 {
        self.radius
    }

    /// Compute the magnetic field at a sensor location due to a dipole
    ///
    /// Uses the Sarvas formula for a current dipole in a spherical conductor.
    ///
    /// # Arguments
    /// * `dipole_pos` - Dipole position in meters
    /// * `dipole_moment` - Dipole moment (current * length) in A*m
    /// * `sensor_pos` - Sensor position in meters
    ///
    /// # Returns
    /// Magnetic field vector [Bx, By, Bz] in Tesla
    pub fn compute_field(
        &self,
        dipole_pos: &Position,
        dipole_moment: &Vector3<f64>,
        sensor_pos: &Position,
    ) -> ForwardResult<Vector3<f64>> {
        // Convert to sphere-centered coordinates
        let r_q = dipole_pos - self.origin; // dipole position relative to sphere
        let r_p = sensor_pos - self.origin; // sensor position relative to sphere

        let r_q_norm = norm(&r_q);

        // Check if dipole is inside the sphere
        if r_q_norm >= self.radius {
            return Err(ForwardError::SourceOutsideHead(format!(
                "Dipole at distance {:.4} m is outside sphere of radius {:.4} m",
                r_q_norm, self.radius
            )));
        }

        // Compute Sarvas formula components
        let a = r_p - r_q; // vector from dipole to sensor
        let a_norm = norm(&a);
        let r_p_norm = norm(&r_p);

        if a_norm < 1e-15 || r_p_norm < 1e-15 {
            return Ok(Vector3::zeros());
        }

        // F = a * (r_p * a + r_p^2 - r_q . r_p)
        let f_scalar = a_norm * (r_p_norm * a_norm + r_p_norm * r_p_norm - dot(&r_q, &r_p));

        if f_scalar.abs() < 1e-30 {
            return Ok(Vector3::zeros());
        }

        // grad_F = (a^2/r_p + a.r_p/a + 2*a + 2*r_p) * r_p - (a + 2*r_p + a.r_p/a) * r_q
        let a_dot_rp = dot(&a, &r_p);
        let term1 = a_norm * a_norm / r_p_norm + a_dot_rp / a_norm + 2.0 * a_norm + 2.0 * r_p_norm;
        let term2 = a_norm + 2.0 * r_p_norm + a_dot_rp / a_norm;
        let grad_f = r_p * term1 - r_q * term2;

        // B = (mu_0 / 4*pi) * (F * (Q x r_q) - (Q x r_q . r_p) * grad_F) / F^2
        let q_cross_rq = cross(dipole_moment, &r_q);
        let q_cross_rq_dot_rp = dot(&q_cross_rq, &r_p);

        let numerator = q_cross_rq * f_scalar - grad_f * q_cross_rq_dot_rp;
        let field = numerator * (MU_0 / (4.0 * PI * f_scalar * f_scalar));

        Ok(field)
    }

    /// Compute the radial component of field for a gradiometer coil
    pub fn compute_coil_field(
        &self,
        dipole_pos: &Position,
        dipole_moment: &Vector3<f64>,
        coil: &MegCoil,
    ) -> ForwardResult<f64> {
        // Compute field at each integration point and sum
        let mut total_flux = 0.0;

        for (pos, weight) in coil.integration_points() {
            let field = self.compute_field(dipole_pos, dipole_moment, pos)?;
            // Project field onto coil normal and weight
            let flux = dot(&field, coil.orientation()) * weight;
            total_flux += flux;
        }

        Ok(total_flux)
    }

    /// Compute the gain matrix (lead field matrix) for all sources and sensors
    ///
    /// # Arguments
    /// * `sources` - Source space with dipole locations and orientations
    /// * `sensors` - Sensor array with MEG sensors
    ///
    /// # Returns
    /// Gain matrix [n_sensors x (3 * n_sources)] for free orientation
    /// or [n_sensors x n_sources] for fixed orientation
    pub fn compute_gain(
        &self,
        sources: &SourceSpace,
        sensors: &SensorArray,
    ) -> ForwardResult<Vec<Vec<f64>>> {
        let n_sensors = sensors.len();
        let n_sources = sources.len();

        // For fixed orientation sources
        if sources.is_fixed_orientation() {
            let mut gain = vec![vec![0.0; n_sources]; n_sensors];

            for (s_idx, sensor) in sensors.iter().enumerate() {
                for (src_idx, source) in sources.iter().enumerate() {
                    let dipole_pos = source.position();
                    let dipole_ori = source
                        .orientation()
                        .unwrap_or_else(|| Vector3::new(0.0, 0.0, 1.0));

                    // Compute field for unit dipole
                    let field = match sensor {
                        Sensor::Magnetometer {
                            position,
                            orientation,
                            ..
                        } => {
                            let b = self.compute_field(dipole_pos, &dipole_ori, position)?;
                            dot(&b, orientation)
                        }
                        Sensor::Gradiometer { coil, .. } => {
                            self.compute_coil_field(dipole_pos, &dipole_ori, coil)?
                        }
                    };

                    gain[s_idx][src_idx] = field;
                }
            }

            Ok(gain)
        } else {
            // Free orientation: 3 columns per source (x, y, z)
            let mut gain = vec![vec![0.0; 3 * n_sources]; n_sensors];

            for (s_idx, sensor) in sensors.iter().enumerate() {
                for (src_idx, source) in sources.iter().enumerate() {
                    let dipole_pos = source.position();

                    // Compute field for each orientation (x, y, z)
                    for (ori_idx, dipole_ori) in [
                        Vector3::new(1.0, 0.0, 0.0),
                        Vector3::new(0.0, 1.0, 0.0),
                        Vector3::new(0.0, 0.0, 1.0),
                    ]
                    .iter()
                    .enumerate()
                    {
                        let field = match sensor {
                            Sensor::Magnetometer {
                                position,
                                orientation,
                                ..
                            } => {
                                let b = self.compute_field(dipole_pos, dipole_ori, position)?;
                                dot(&b, orientation)
                            }
                            Sensor::Gradiometer { coil, .. } => {
                                self.compute_coil_field(dipole_pos, dipole_ori, coil)?
                            }
                        };

                        gain[s_idx][3 * src_idx + ori_idx] = field;
                    }
                }
            }

            Ok(gain)
        }
    }

    /// Compute gain matrix in parallel using rayon
    pub fn compute_gain_parallel(
        &self,
        sources: &SourceSpace,
        sensors: &SensorArray,
    ) -> ForwardResult<Vec<Vec<f64>>> {
        let n_sources = sources.len();

        if sources.is_fixed_orientation() {
            let gain: Vec<Vec<f64>> = sensors
                .iter()
                .map(|sensor| {
                    sources
                        .iter()
                        .map(|source| {
                            let dipole_pos = source.position();
                            let dipole_ori = source
                                .orientation()
                                .unwrap_or_else(|| Vector3::new(0.0, 0.0, 1.0));

                            match sensor {
                                Sensor::Magnetometer {
                                    position,
                                    orientation,
                                    ..
                                } => {
                                    let b = self
                                        .compute_field(dipole_pos, &dipole_ori, position)
                                        .unwrap_or_else(|_| Vector3::zeros());
                                    dot(&b, orientation)
                                }
                                Sensor::Gradiometer { coil, .. } => self
                                    .compute_coil_field(dipole_pos, &dipole_ori, coil)
                                    .unwrap_or(0.0),
                            }
                        })
                        .collect()
                })
                .collect();

            Ok(gain)
        } else {
            let gain: Vec<Vec<f64>> = sensors
                .iter()
                .map(|sensor| {
                    let mut row = vec![0.0; 3 * n_sources];
                    for (src_idx, source) in sources.iter().enumerate() {
                        let dipole_pos = source.position();

                        for (ori_idx, dipole_ori) in [
                            Vector3::new(1.0, 0.0, 0.0),
                            Vector3::new(0.0, 1.0, 0.0),
                            Vector3::new(0.0, 0.0, 1.0),
                        ]
                        .iter()
                        .enumerate()
                        {
                            let field = match sensor {
                                Sensor::Magnetometer {
                                    position,
                                    orientation,
                                    ..
                                } => {
                                    let b = self
                                        .compute_field(dipole_pos, dipole_ori, position)
                                        .unwrap_or_else(|_| Vector3::zeros());
                                    dot(&b, orientation)
                                }
                                Sensor::Gradiometer { coil, .. } => self
                                    .compute_coil_field(dipole_pos, dipole_ori, coil)
                                    .unwrap_or(0.0),
                            };

                            row[3 * src_idx + ori_idx] = field;
                        }
                    }
                    row
                })
                .collect();

            Ok(gain)
        }
    }
}

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

    #[test]
    fn test_sphere_creation() {
        let sphere = SphericalMeg::new([0.0, 0.0, 0.04], 0.08).unwrap();
        assert_eq!(sphere.radius(), 0.08);
    }

    #[test]
    fn test_invalid_radius() {
        let result = SphericalMeg::new([0.0, 0.0, 0.0], -0.1);
        assert!(result.is_err());
    }

    #[test]
    fn test_field_radial_dipole() {
        // A radially-oriented dipole produces no magnetic field (Sarvas)
        let sphere = SphericalMeg::default_head();
        let dipole_pos = Vector3::new(0.0, 0.0, 0.06); // 2cm inside sphere
        let radial_ori = normalize(&(dipole_pos - sphere.origin()));
        let dipole_moment = radial_ori * 1e-9; // 1 nAm

        let sensor_pos = Vector3::new(0.0, 0.1, 0.08); // outside sphere

        let field = sphere
            .compute_field(&dipole_pos, &dipole_moment, &sensor_pos)
            .unwrap();

        // Radial dipole should produce essentially zero field
        assert!(norm(&field) < 1e-20);
    }

    #[test]
    fn test_field_tangential_dipole() {
        // A tangential dipole should produce non-zero field
        let sphere = SphericalMeg::default_head();
        let dipole_pos = Vector3::new(0.0, 0.0, 0.08); // at sphere center + 4cm
        let dipole_moment = Vector3::new(1e-9, 0.0, 0.0); // x-oriented, 1 nAm

        let sensor_pos = Vector3::new(0.0, 0.1, 0.08); // 10 cm from center

        let field = sphere
            .compute_field(&dipole_pos, &dipole_moment, &sensor_pos)
            .unwrap();

        // Should have non-zero z-component due to tangential dipole
        assert!(norm(&field) > 1e-20);
    }

    #[test]
    fn test_source_outside_sphere() {
        let sphere = SphericalMeg::new([0.0, 0.0, 0.0], 0.08).unwrap();
        let dipole_pos = Vector3::new(0.0, 0.0, 0.1); // outside sphere
        let dipole_moment = Vector3::new(1e-9, 0.0, 0.0);
        let sensor_pos = Vector3::new(0.0, 0.1, 0.0);

        let result = sphere.compute_field(&dipole_pos, &dipole_moment, &sensor_pos);
        assert!(result.is_err());
    }
}