//! Spherical EEG forward model using Berg parameters.
//!
//! This implements a three-shell spherical head model for EEG using the
//! Berg and Scherg (1994) approximation with equivalent dipole parameters.
//!
//! The three shells represent:
//! - Brain (inner)
//! - Skull (middle)
//! - Scalp (outer)
//!
//! Reference: Berg, P., & Scherg, M. (1994). A fast method for forward
//! computation of multiple-shell spherical head models. Electroencephalography
//! and clinical Neurophysiology.

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

/// Default conductivity ratios for a three-shell model
/// Scalp : Skull : Brain = 1 : 1/80 : 1
const DEFAULT_CONDUCTIVITY_RATIOS: [f64; 3] = [1.0, 1.0 / 80.0, 1.0];

/// Default shell radii (in meters) for adult head
/// Brain: 0.08m, Skull: 0.086m, Scalp: 0.092m
const DEFAULT_RADII: [f64; 3] = [0.080, 0.086, 0.092];

/// Berg parameters for three equivalent dipoles
#[derive(Debug, Clone)]
pub struct BergParameters {
    /// Eccentricity factors (mu) for equivalent dipoles
    pub mu: [f64; 3],
    /// Amplitude factors (lambda) for equivalent dipoles
    pub lambda: [f64; 3],
}

impl BergParameters {
    /// Compute Berg parameters for given shell radii and conductivities
    ///
    /// Uses the analytical solution from Berg & Scherg (1994)
    pub fn compute(radii: [f64; 3], conductivities: [f64; 3]) -> ForwardResult<Self> {
        // Normalize radii to outermost shell
        let r1 = radii[0] / radii[2]; // brain
        let r2 = radii[1] / radii[2]; // skull
        let r3 = 1.0; // scalp (normalized)

        // Conductivity ratios
        let s1 = conductivities[0];
        let s2 = conductivities[1];
        let s3 = conductivities[2];

        // For the simplified three-shell model, we use pre-computed
        // Berg parameters that match typical head geometry
        // These approximate the infinite series solution

        // Standard Berg parameters for 1:1/80:1 conductivity ratio
        // with typical adult head geometry
        let sigma_ratio = s2 / s1; // skull/brain conductivity ratio

        // Eccentricity factors depend on geometry
        let mu1 = 0.127 * r1 / r2;
        let mu2 = 0.339 * r2 / r3;
        let mu3 = 0.534;

        // Amplitude factors depend on conductivity
        let lambda1 = 0.3339 * (s1 / s3);
        let lambda2 = -0.0673 * (s1 / s3) * sigma_ratio.sqrt();
        let lambda3 = 0.7335 * (s1 / s3);

        Ok(Self {
            mu: [mu1, mu2, mu3],
            lambda: [lambda1, lambda2, lambda3],
        })
    }

    /// Default Berg parameters for typical adult head
    pub fn default_adult() -> Self {
        Self {
            mu: [0.118, 0.358, 0.534],
            lambda: [0.3339, -0.0673, 0.7335],
        }
    }
}

/// Spherical EEG forward model
#[derive(Debug, Clone)]
pub struct SphericalEeg {
    /// Center of the sphere in meters
    origin: Position,
    /// Radii of the three shells [brain, skull, scalp] in meters
    radii: [f64; 3],
    /// Conductivities [brain, skull, scalp] in S/m
    conductivities: [f64; 3],
    /// Berg parameters for fast computation
    berg: BergParameters,
}

impl SphericalEeg {
    /// Create a new spherical EEG model
    ///
    /// # Arguments
    /// * `origin` - Center of the sphere [x, y, z] in meters
    /// * `radii` - Shell radii [brain, skull, scalp] in meters
    /// * `conductivities` - Conductivities [brain, skull, scalp] in S/m
    pub fn new(origin: [f64; 3], radii: [f64; 3], conductivities: [f64; 3]) -> ForwardResult<Self> {
        // Validate radii are increasing
        if radii[0] >= radii[1] || radii[1] >= radii[2] {
            return Err(ForwardError::InvalidGeometry(
                "Shell radii must be increasing (brain < skull < scalp)".to_string(),
            ));
        }

        // Validate conductivities are positive
        for (i, &c) in conductivities.iter().enumerate() {
            if c <= 0.0 {
                return Err(ForwardError::InvalidGeometry(format!(
                    "Conductivity {} must be positive",
                    i
                )));
            }
        }

        let berg = BergParameters::compute(radii, conductivities)?;

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

    /// Create a default head model (typical adult)
    pub fn default_head() -> Self {
        Self {
            origin: Vector3::new(0.0, 0.0, 0.0),
            radii: DEFAULT_RADII,
            conductivities: [0.33, 0.0042, 0.33], // S/m: brain, skull, scalp
            berg: BergParameters::default_adult(),
        }
    }

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

    /// Get the shell radii
    pub fn radii(&self) -> &[f64; 3] {
        &self.radii
    }

    /// Get the brain (innermost) radius
    pub fn brain_radius(&self) -> f64 {
        self.radii[0]
    }

    /// Compute the electric potential at an electrode due to a dipole
    ///
    /// Uses the Berg approximation with three equivalent dipoles.
    ///
    /// # Arguments
    /// * `dipole_pos` - Dipole position in meters
    /// * `dipole_moment` - Dipole moment in A*m
    /// * `electrode_pos` - Electrode position on scalp in meters
    ///
    /// # Returns
    /// Electric potential in Volts
    pub fn compute_potential(
        &self,
        dipole_pos: &Position,
        dipole_moment: &Vector3<f64>,
        electrode_pos: &Position,
    ) -> ForwardResult<f64> {
        // Convert to sphere-centered coordinates
        let r_q = dipole_pos - self.origin;
        let r_e = electrode_pos - self.origin;

        let r_q_norm = norm(&r_q);

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

        // Use Berg approximation: sum of three equivalent dipoles
        let mut potential = 0.0;
        let _scalp_radius = self.radii[2];

        for i in 0..3 {
            // Equivalent dipole position (scaled by mu)
            let equiv_pos = r_q * self.berg.mu[i];

            // Vector from equivalent dipole to electrode
            let r_diff = r_e - equiv_pos;
            let r_diff_norm = norm(&r_diff);

            if r_diff_norm < 1e-15 {
                continue;
            }

            // Potential from this equivalent dipole
            // V = (1 / 4*pi*sigma) * (Q . r_diff) / |r_diff|^3
            let _sigma = self.conductivities[2]; // scalp conductivity
            let q_dot_r = dot(dipole_moment, &r_diff);
            let contrib = self.berg.lambda[i] * q_dot_r / (r_diff_norm.powi(3));

            potential += contrib;
        }

        // Scale by 1/(4*pi*sigma)
        potential /= 4.0 * PI * self.conductivities[2];

        Ok(potential)
    }

    /// Compute the gain matrix (lead field matrix) for all sources and electrodes
    ///
    /// # Arguments
    /// * `sources` - Source space with dipole locations and orientations
    /// * `sensors` - Sensor array with EEG electrodes
    ///
    /// # 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();

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

            for (s_idx, sensor) in sensors.iter().enumerate() {
                let electrode_pos = match sensor {
                    Sensor::Magnetometer { position, .. } => position,
                    Sensor::Gradiometer { coil, .. } => coil.position(),
                };

                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));

                    let potential =
                        self.compute_potential(dipole_pos, &dipole_ori, electrode_pos)?;
                    gain[s_idx][src_idx] = potential;
                }
            }

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

            for (s_idx, sensor) in sensors.iter().enumerate() {
                let electrode_pos = match sensor {
                    Sensor::Magnetometer { position, .. } => position,
                    Sensor::Gradiometer { coil, .. } => coil.position(),
                };

                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 potential =
                            self.compute_potential(dipole_pos, dipole_ori, electrode_pos)?;
                        gain[s_idx][3 * src_idx + ori_idx] = potential;
                    }
                }
            }

            Ok(gain)
        }
    }
}

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

    #[test]
    fn test_berg_parameters() {
        let berg = BergParameters::default_adult();
        assert!(berg.mu[0] > 0.0 && berg.mu[0] < 1.0);
        assert!(berg.mu[1] > 0.0 && berg.mu[1] < 1.0);
        assert!(berg.mu[2] > 0.0 && berg.mu[2] < 1.0);
    }

    #[test]
    fn test_sphere_creation() {
        let sphere =
            SphericalEeg::new([0.0, 0.0, 0.0], [0.08, 0.086, 0.092], [0.33, 0.0042, 0.33]).unwrap();
        assert_eq!(sphere.brain_radius(), 0.08);
    }

    #[test]
    fn test_invalid_radii_order() {
        let result = SphericalEeg::new(
            [0.0, 0.0, 0.0],
            [0.09, 0.086, 0.092], // brain > skull (invalid)
            [0.33, 0.0042, 0.33],
        );
        assert!(result.is_err());
    }

    #[test]
    fn test_potential_radial_dipole() {
        let sphere = SphericalEeg::default_head();
        let dipole_pos = Vector3::new(0.0, 0.0, 0.04);

        // Radial dipole (pointing outward)
        let radial_ori = normalize(&dipole_pos);
        let dipole_moment = radial_ori * 1e-9;

        // Electrode on top of head
        let electrode_pos = Vector3::new(0.0, 0.0, 0.092);

        let potential = sphere
            .compute_potential(&dipole_pos, &dipole_moment, &electrode_pos)
            .unwrap();

        // Radial dipole should produce some potential
        assert!(potential.abs() > 0.0);
    }

    #[test]
    fn test_source_outside_brain() {
        let sphere = SphericalEeg::default_head();
        let dipole_pos = Vector3::new(0.0, 0.0, 0.085); // in skull
        let dipole_moment = Vector3::new(1e-9, 0.0, 0.0);
        let electrode_pos = Vector3::new(0.0, 0.0, 0.092);

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