//! eLORETA (Exact Low Resolution Electromagnetic Tomography) inverse solution.
//!
//! eLORETA is an inverse method with the unique property of exact (zero-error)
//! localization for point sources. Unlike sLORETA which only standardizes the
//! current density, eLORETA computes weights that ensure the resolution matrix
//! has specific properties guaranteeing zero localization error.
//!
//! ## Mathematical Background
//!
//! The eLORETA solution is: J = W * M
//!
//! where W is computed iteratively such that the resulting resolution matrix
//! R = W * G satisfies specific properties.
//!
//! The weight matrix W_i for source i is:
//! W_i = [G_i^T * C^(-1) * G_i]^(-1/2) for each source orientation block
//!
//! This ensures that the point-spread function is centered on the true source.
//!
//! ## References
//!
//! - Pascual-Marqui, R.D. (2007). Discrete, 3D distributed, linear imaging methods
//!   of electric neuronal activity. Part 1: exact, zero error localization.
//! - Pascual-Marqui, R.D. (2009). Theory of the EEG inverse problem.
//!
//! ## Usage
//!
//! ```rust,ignore
//! use rtx_neuro_inverse::loreta::{EloretaInverse, EloretaConfig};
//!
//! let config = EloretaConfig::default();
//! let inverse = EloretaInverse::make_inverse(&gain, &noise_cov, config)?;
//! let stc = inverse.apply(&evoked_data)?;
//! ```

use crate::{Covariance, InverseError, InverseResult, SourceEstimate};
use nalgebra::{DMatrix, DVector, Matrix3};
use rtx_neuro_forward::GainMatrix;

/// Configuration for eLORETA inverse
#[derive(Debug, Clone)]
pub struct EloretaConfig {
    /// Regularization parameter (lambda^2)
    pub lambda2: f64,
    /// Maximum iterations for weight computation
    pub max_iter: usize,
    /// Convergence tolerance
    pub tol: f64,
    /// Depth weighting exponent (0 = none, typically 0.5-0.9)
    pub depth: f64,
    /// Whether sources have free orientation
    pub free_orientation: bool,
}

impl Default for EloretaConfig {
    fn default() -> Self {
        Self {
            lambda2: 1.0 / 9.0, // SNR^2 = 9
            max_iter: 100,
            tol: 1e-6,
            depth: 0.5,
            free_orientation: true,
        }
    }
}

impl EloretaConfig {
    /// Create config for fixed orientation sources
    pub fn fixed() -> Self {
        Self {
            free_orientation: false,
            ..Self::default()
        }
    }

    /// Set regularization parameter
    pub fn with_lambda2(mut self, lambda2: f64) -> Self {
        self.lambda2 = lambda2;
        self
    }

    /// Set depth weighting
    pub fn with_depth(mut self, depth: f64) -> Self {
        self.depth = depth;
        self
    }
}

/// eLORETA inverse operator
#[derive(Debug, Clone)]
pub struct EloretaInverse {
    /// Inverse kernel [n_source_columns x n_channels]
    kernel: DMatrix<f64>,
    /// Weight matrix for each source (for normalization)
    weights: Vec<DMatrix<f64>>,
    /// Number of sources
    n_sources: usize,
    /// Number of channels
    n_channels: usize,
    /// Whether sources have free orientation
    free_orientation: bool,
    /// Configuration used
    config: EloretaConfig,
    /// Source indices
    source_indices: Vec<usize>,
}

impl EloretaInverse {
    /// Create an eLORETA inverse operator
    ///
    /// # Arguments
    /// * `gain` - Forward model gain matrix [n_channels x n_sources]
    /// * `noise_cov` - Noise covariance matrix
    /// * `config` - eLORETA configuration
    pub fn make_inverse(
        gain: &GainMatrix,
        noise_cov: &Covariance,
        config: EloretaConfig,
    ) -> InverseResult<Self> {
        let n_channels = gain.n_sensors();
        let n_source_cols = gain.n_source_columns();
        let free_orientation = gain.is_free_orientation();

        if noise_cov.n_channels() != n_channels {
            return Err(InverseError::DimensionMismatch(format!(
                "Noise covariance has {} channels, gain has {}",
                noise_cov.n_channels(),
                n_channels
            )));
        }

        // Convert gain to nalgebra matrix
        let g = Self::gain_to_matrix(gain);

        // Compute inverse of regularized noise covariance
        let c_inv = Self::compute_regularized_inv(noise_cov, config.lambda2)?;

        // Compute depth weights
        let depth_weights = Self::compute_depth_weights(&g, config.depth, free_orientation);

        // Compute eLORETA weights iteratively
        let (weights, kernel) = if free_orientation {
            Self::compute_eloreta_weights_free(&g, &c_inv, &depth_weights, &config)?
        } else {
            Self::compute_eloreta_weights_fixed(&g, &c_inv, &depth_weights, &config)?
        };

        let n_sources = if free_orientation {
            n_source_cols / 3
        } else {
            n_source_cols
        };

        Ok(Self {
            kernel,
            weights,
            n_sources,
            n_channels,
            free_orientation,
            config,
            source_indices: (0..n_sources).collect(),
        })
    }

    /// Apply the inverse operator to sensor data
    ///
    /// # Arguments
    /// * `data` - Sensor data [n_channels x n_times]
    ///
    /// # Returns
    /// Source estimates
    pub fn apply(&self, data: &[Vec<f64>]) -> InverseResult<SourceEstimate> {
        if data.len() != self.n_channels {
            return Err(InverseError::DimensionMismatch(format!(
                "Expected {} channels, got {}",
                self.n_channels,
                data.len()
            )));
        }

        let n_times = data[0].len();

        // Convert to matrix
        let data_mat = DMatrix::from_fn(self.n_channels, n_times, |i, j| data[i][j]);

        // Apply inverse kernel
        let source_mat = &self.kernel * &data_mat;

        // Convert to output format
        let n_rows = source_mat.nrows();
        let source_data: Vec<Vec<f64>> = (0..n_rows)
            .map(|i| (0..n_times).map(|j| source_mat[(i, j)]).collect())
            .collect();

        let times: Vec<f64> = (0..n_times).map(|i| i as f64).collect();

        Ok(SourceEstimate::new(
            source_data,
            times,
            self.source_indices.clone(),
            self.free_orientation,
        ))
    }

    /// Get the inverse kernel matrix
    pub fn kernel(&self) -> &DMatrix<f64> {
        &self.kernel
    }

    /// Get number of sources
    pub fn n_sources(&self) -> usize {
        self.n_sources
    }

    /// Get the weight matrices
    pub fn weights(&self) -> &[DMatrix<f64>] {
        &self.weights
    }

    // ========== Private methods ==========

    /// Convert GainMatrix to nalgebra DMatrix
    fn gain_to_matrix(gain: &GainMatrix) -> DMatrix<f64> {
        let n_rows = gain.n_sensors();
        let n_cols = gain.n_source_columns();
        let data = gain.data();
        DMatrix::from_fn(n_rows, n_cols, |i, j| data[i][j])
    }

    /// Compute regularized inverse of noise covariance
    fn compute_regularized_inv(
        noise_cov: &Covariance,
        lambda2: f64,
    ) -> InverseResult<DMatrix<f64>> {
        let n = noise_cov.n_channels();
        let c = noise_cov.data();

        // Add regularization: C_reg = C + lambda^2 * trace(C)/n * I
        let trace: f64 = (0..n).map(|i| c[(i, i)]).sum();
        let reg_val = lambda2 * trace / n as f64;

        let mut c_reg = c.clone();
        for i in 0..n {
            c_reg[(i, i)] += reg_val;
        }

        // Compute inverse via eigendecomposition for stability
        let eigen = c_reg.symmetric_eigen();
        let eigenvalues = eigen.eigenvalues;
        let eigenvectors = eigen.eigenvectors;

        let min_eig = eigenvalues.iter().copied().fold(f64::INFINITY, f64::min);
        if min_eig <= 0.0 {
            return Err(InverseError::ComputationError(format!(
                "Regularized covariance has non-positive eigenvalue: {:.2e}",
                min_eig
            )));
        }

        let d_inv = DMatrix::from_diagonal(&DVector::from_fn(n, |i, _| 1.0 / eigenvalues[i]));

        Ok(&eigenvectors * &d_inv * eigenvectors.transpose())
    }

    /// Compute depth weights
    fn compute_depth_weights(gain: &DMatrix<f64>, depth: f64, free_orientation: bool) -> Vec<f64> {
        let n_cols = gain.ncols();

        if depth == 0.0 {
            return vec![1.0; n_cols];
        }

        if free_orientation {
            let n_sources = n_cols / 3;
            let mut weights = Vec::with_capacity(n_cols);

            for src in 0..n_sources {
                // Compute norm across all orientations
                let mut sum_sq = 0.0;
                for ori in 0..3 {
                    let col_idx = 3 * src + ori;
                    for row in 0..gain.nrows() {
                        sum_sq += gain[(row, col_idx)].powi(2);
                    }
                }
                let norm = sum_sq.sqrt();
                let w = if norm > 1e-15 { norm.powf(-depth) } else { 1.0 };

                // Same weight for all orientations
                weights.push(w);
                weights.push(w);
                weights.push(w);
            }
            weights
        } else {
            (0..n_cols)
                .map(|col| {
                    let sum_sq: f64 = (0..gain.nrows()).map(|row| gain[(row, col)].powi(2)).sum();
                    let norm = sum_sq.sqrt();
                    if norm > 1e-15 { norm.powf(-depth) } else { 1.0 }
                })
                .collect()
        }
    }

    /// Compute eLORETA weights for fixed orientation sources
    fn compute_eloreta_weights_fixed(
        gain: &DMatrix<f64>,
        c_inv: &DMatrix<f64>,
        depth_weights: &[f64],
        config: &EloretaConfig,
    ) -> InverseResult<(Vec<DMatrix<f64>>, DMatrix<f64>)> {
        let n_channels = gain.nrows();
        let n_sources = gain.ncols();

        // Initialize weights to identity (scaled by depth)
        let mut source_weights: Vec<f64> = depth_weights.to_vec();

        // Iterative weight computation
        for _iter in 0..config.max_iter {
            let old_weights = source_weights.clone();

            // Compute H = sum_i (w_i^2 * g_i * g_i^T)
            let mut h = DMatrix::zeros(n_channels, n_channels);
            for i in 0..n_sources {
                let g_i = gain.column(i);
                let w_sq = source_weights[i].powi(2);
                for row in 0..n_channels {
                    for col in 0..n_channels {
                        h[(row, col)] += w_sq * g_i[row] * g_i[col];
                    }
                }
            }

            // Regularize H
            let trace: f64 = (0..n_channels).map(|i| h[(i, i)]).sum();
            let reg = config.lambda2 * trace / n_channels as f64;
            for i in 0..n_channels {
                h[(i, i)] += reg;
            }

            // Compute T = (H + λC)^(-1)
            let h_reg = &h + c_inv * config.lambda2;
            let t = match h_reg.clone().try_inverse() {
                Some(inv) => inv,
                None => {
                    return Err(InverseError::ComputationError(
                        "Failed to invert H matrix".to_string(),
                    ));
                }
            };

            // Update weights: w_i = (g_i^T * T * g_i)^(-1/4)
            for i in 0..n_sources {
                let g_i = gain.column(i);
                let tg = &t * g_i;
                let gtg = g_i.dot(&tg);

                if gtg > 1e-30 {
                    source_weights[i] = depth_weights[i] * gtg.powf(-0.25);
                }
            }

            // Check convergence
            let change: f64 = source_weights
                .iter()
                .zip(old_weights.iter())
                .map(|(a, b)| (a - b).abs())
                .sum::<f64>()
                / n_sources as f64;

            if change < config.tol {
                break;
            }
        }

        // Compute final kernel: K = W * G^T * T
        // where W is diagonal weight matrix
        let mut h = DMatrix::zeros(n_channels, n_channels);
        for i in 0..n_sources {
            let g_i = gain.column(i);
            let w_sq = source_weights[i].powi(2);
            for row in 0..n_channels {
                for col in 0..n_channels {
                    h[(row, col)] += w_sq * g_i[row] * g_i[col];
                }
            }
        }

        let trace: f64 = (0..n_channels).map(|i| h[(i, i)]).sum();
        let reg = config.lambda2 * trace / n_channels as f64;
        for i in 0..n_channels {
            h[(i, i)] += reg;
        }

        let h_reg = &h + c_inv * config.lambda2;
        let t = h_reg.try_inverse().ok_or_else(|| {
            InverseError::ComputationError("Failed to compute final inverse".to_string())
        })?;

        // Kernel [n_sources x n_channels]
        let mut kernel = DMatrix::zeros(n_sources, n_channels);
        for i in 0..n_sources {
            let g_i = gain.column(i);
            let w_sq = source_weights[i].powi(2);
            let tg = &t * g_i;
            for ch in 0..n_channels {
                kernel[(i, ch)] = w_sq * tg[ch];
            }
        }

        // Convert weights to matrix format
        let weight_matrices: Vec<DMatrix<f64>> = source_weights
            .iter()
            .map(|&w| DMatrix::from_element(1, 1, w))
            .collect();

        Ok((weight_matrices, kernel))
    }

    /// Compute eLORETA weights for free orientation sources
    fn compute_eloreta_weights_free(
        gain: &DMatrix<f64>,
        c_inv: &DMatrix<f64>,
        depth_weights: &[f64],
        config: &EloretaConfig,
    ) -> InverseResult<(Vec<DMatrix<f64>>, DMatrix<f64>)> {
        let n_channels = gain.nrows();
        let n_sources = gain.ncols() / 3;

        // Initialize weight matrices (3x3 for each source)
        let mut source_weights: Vec<Matrix3<f64>> = (0..n_sources)
            .map(|i| {
                let d = depth_weights[3 * i];
                Matrix3::identity() * d
            })
            .collect();

        // Iterative weight computation
        for _iter in 0..config.max_iter {
            let old_weights: Vec<Matrix3<f64>> = source_weights.clone();

            // Compute H = sum_i (G_i * W_i^2 * G_i^T)
            let mut h = DMatrix::zeros(n_channels, n_channels);
            for i in 0..n_sources {
                let g_i = Self::extract_source_gain(gain, i);
                let w_sq = source_weights[i] * source_weights[i];

                // H += G_i * W_i^2 * G_i^T
                for ch1 in 0..n_channels {
                    for ch2 in 0..n_channels {
                        let mut sum = 0.0;
                        for ori1 in 0..3 {
                            for ori2 in 0..3 {
                                sum += g_i[(ch1, ori1)] * w_sq[(ori1, ori2)] * g_i[(ch2, ori2)];
                            }
                        }
                        h[(ch1, ch2)] += sum;
                    }
                }
            }

            // Regularize H
            let trace: f64 = (0..n_channels).map(|i| h[(i, i)]).sum();
            let reg = config.lambda2 * trace / n_channels as f64;
            for i in 0..n_channels {
                h[(i, i)] += reg;
            }

            // Compute T = (H + λC)^(-1)
            let h_reg = &h + c_inv * config.lambda2;
            let t = match h_reg.clone().try_inverse() {
                Some(inv) => inv,
                None => {
                    return Err(InverseError::ComputationError(
                        "Failed to invert H matrix".to_string(),
                    ));
                }
            };

            // Update weights: W_i = D_i * (G_i^T * T * G_i)^(-1/4)
            for i in 0..n_sources {
                let g_i = Self::extract_source_gain(gain, i);

                // Compute G_i^T * T * G_i (3x3 matrix)
                let mut gtg = Matrix3::zeros();
                for ori1 in 0..3 {
                    for ori2 in 0..3 {
                        let mut sum = 0.0;
                        for ch1 in 0..n_channels {
                            for ch2 in 0..n_channels {
                                sum += g_i[(ch1, ori1)] * t[(ch1, ch2)] * g_i[(ch2, ori2)];
                            }
                        }
                        gtg[(ori1, ori2)] = sum;
                    }
                }

                // Compute matrix power -1/4 via eigendecomposition
                if let Some(new_w) = Self::matrix_power_quarter_inv(&gtg) {
                    let d = depth_weights[3 * i];
                    source_weights[i] = new_w * d;
                }
            }

            // Check convergence
            let change: f64 = source_weights
                .iter()
                .zip(old_weights.iter())
                .map(|(a, b)| (a - b).norm())
                .sum::<f64>()
                / n_sources as f64;

            if change < config.tol {
                break;
            }
        }

        // Compute final kernel
        let mut h = DMatrix::zeros(n_channels, n_channels);
        for i in 0..n_sources {
            let g_i = Self::extract_source_gain(gain, i);
            let w_sq = source_weights[i] * source_weights[i];

            for ch1 in 0..n_channels {
                for ch2 in 0..n_channels {
                    let mut sum = 0.0;
                    for ori1 in 0..3 {
                        for ori2 in 0..3 {
                            sum += g_i[(ch1, ori1)] * w_sq[(ori1, ori2)] * g_i[(ch2, ori2)];
                        }
                    }
                    h[(ch1, ch2)] += sum;
                }
            }
        }

        let trace: f64 = (0..n_channels).map(|i| h[(i, i)]).sum();
        let reg = config.lambda2 * trace / n_channels as f64;
        for i in 0..n_channels {
            h[(i, i)] += reg;
        }

        let h_reg = &h + c_inv * config.lambda2;
        let t = h_reg.try_inverse().ok_or_else(|| {
            InverseError::ComputationError("Failed to compute final inverse".to_string())
        })?;

        // Kernel [3*n_sources x n_channels]
        let mut kernel = DMatrix::zeros(3 * n_sources, n_channels);
        for i in 0..n_sources {
            let g_i = Self::extract_source_gain(gain, i);
            let w_sq = source_weights[i] * source_weights[i];

            for ori in 0..3 {
                for ch in 0..n_channels {
                    let mut sum = 0.0;
                    for ori2 in 0..3 {
                        for ch2 in 0..n_channels {
                            sum += w_sq[(ori, ori2)] * g_i[(ch2, ori2)] * t[(ch2, ch)];
                        }
                    }
                    kernel[(3 * i + ori, ch)] = sum;
                }
            }
        }

        // Convert to output format
        let weight_matrices: Vec<DMatrix<f64>> = source_weights
            .iter()
            .map(|w| {
                let mut m = DMatrix::zeros(3, 3);
                for i in 0..3 {
                    for j in 0..3 {
                        m[(i, j)] = w[(i, j)];
                    }
                }
                m
            })
            .collect();

        Ok((weight_matrices, kernel))
    }

    /// Extract 3-column gain matrix for a single source
    fn extract_source_gain(gain: &DMatrix<f64>, source_idx: usize) -> DMatrix<f64> {
        let n_channels = gain.nrows();
        let start_col = 3 * source_idx;

        DMatrix::from_fn(n_channels, 3, |row, col| gain[(row, start_col + col)])
    }

    /// Compute A^(-1/4) for a 3x3 symmetric positive definite matrix
    fn matrix_power_quarter_inv(a: &Matrix3<f64>) -> Option<Matrix3<f64>> {
        // Convert to nalgebra for eigendecomposition
        let eigen = a.symmetric_eigen();

        // Check all eigenvalues are positive
        for &ev in eigen.eigenvalues.iter() {
            if ev <= 1e-30 {
                return None;
            }
        }

        // Compute D^(-1/4)
        let d_inv_quarter = Matrix3::from_diagonal(&nalgebra::Vector3::new(
            eigen.eigenvalues[0].powf(-0.25),
            eigen.eigenvalues[1].powf(-0.25),
            eigen.eigenvalues[2].powf(-0.25),
        ));

        // A^(-1/4) = V * D^(-1/4) * V^T
        Some(eigen.eigenvectors * d_inv_quarter * eigen.eigenvectors.transpose())
    }
}

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

    fn create_simple_gain() -> GainMatrix {
        // 4 sensors, 2 sources with free orientation (6 columns)
        let data = vec![
            vec![1.0, 0.2, 0.1, 0.3, 0.5, 0.2],
            vec![0.2, 1.0, 0.3, 0.5, 0.3, 0.1],
            vec![0.1, 0.3, 1.0, 0.2, 0.1, 0.5],
            vec![0.3, 0.1, 0.2, 1.0, 0.2, 0.3],
        ];
        let names = vec![
            "S1".to_string(),
            "S2".to_string(),
            "S3".to_string(),
            "S4".to_string(),
        ];
        GainMatrix::new(data, true, names).unwrap()
    }

    fn create_fixed_gain() -> GainMatrix {
        // 4 sensors, 3 fixed sources
        let data = vec![
            vec![1.0, 0.5, 0.2],
            vec![0.5, 1.0, 0.5],
            vec![0.2, 0.5, 1.0],
            vec![0.1, 0.3, 0.6],
        ];
        let names = vec![
            "S1".to_string(),
            "S2".to_string(),
            "S3".to_string(),
            "S4".to_string(),
        ];
        GainMatrix::new(data, false, names).unwrap()
    }

    #[test]
    fn test_eloreta_config() {
        let config = EloretaConfig::default();
        assert!(config.free_orientation);
        assert!((config.lambda2 - 1.0 / 9.0).abs() < 1e-10);

        let config = EloretaConfig::fixed().with_lambda2(0.1);
        assert!(!config.free_orientation);
        assert!((config.lambda2 - 0.1).abs() < 1e-10);
    }

    #[test]
    fn test_make_eloreta_fixed() {
        let gain = create_fixed_gain();
        let noise_cov = Covariance::identity(4, CovarianceType::Noise);

        let config = EloretaConfig::fixed();
        let inv = EloretaInverse::make_inverse(&gain, &noise_cov, config).unwrap();

        assert_eq!(inv.n_sources(), 3);
        assert!(!inv.free_orientation);
    }

    #[test]
    fn test_make_eloreta_free() {
        let gain = create_simple_gain();
        let noise_cov = Covariance::identity(4, CovarianceType::Noise);

        let config = EloretaConfig::default();
        let inv = EloretaInverse::make_inverse(&gain, &noise_cov, config).unwrap();

        assert_eq!(inv.n_sources(), 2); // 6 columns / 3 orientations
        assert!(inv.free_orientation);
    }

    #[test]
    fn test_apply_eloreta() {
        let gain = create_fixed_gain();
        let noise_cov = Covariance::identity(4, CovarianceType::Noise);

        let config = EloretaConfig::fixed();
        let inv = EloretaInverse::make_inverse(&gain, &noise_cov, config).unwrap();

        let data = vec![
            vec![1.0, 0.0],
            vec![0.0, 1.0],
            vec![0.5, 0.5],
            vec![0.2, 0.8],
        ];

        let stc = inv.apply(&data).unwrap();
        assert_eq!(stc.n_sources(), 3);
        assert_eq!(stc.n_times(), 2);
    }

    #[test]
    fn test_eloreta_kernel_shape() {
        let gain = create_simple_gain();
        let noise_cov = Covariance::identity(4, CovarianceType::Noise);

        let inv =
            EloretaInverse::make_inverse(&gain, &noise_cov, EloretaConfig::default()).unwrap();

        let kernel = inv.kernel();
        assert_eq!(kernel.nrows(), 6); // 3 orientations * 2 sources
        assert_eq!(kernel.ncols(), 4); // 4 channels
    }
}