//! Beamformer methods for source localization.
//!
//! This module provides beamformer inverse solutions:
//! - **LCMV** (Linearly Constrained Minimum Variance): Time-domain beamformer
//! - **DICS** (Dynamic Imaging of Coherent Sources): Frequency-domain beamformer
//!
//! ## Mathematical Background
//!
//! The beamformer filter is: w = (C^-1 * g) / (g^T * C^-1 * g)
//! where C is covariance/CSD, g is lead field for the target location.
//!
//! LCMV uses time-domain covariance, while DICS uses cross-spectral density.

use crate::{Covariance, InverseError, InverseResult, SourceEstimate};
use nalgebra::{Complex, DMatrix, DVector};
use rtx_neuro_forward::GainMatrix;
use std::f64::consts::PI;

/// LCMV beamformer
#[derive(Debug, Clone)]
pub struct LcmvBeamformer {
    /// Spatial filters [n_sources x n_channels]
    filters: DMatrix<f64>,
    /// Number of sources
    n_sources: usize,
    /// Number of channels
    n_channels: usize,
    /// Regularization parameter
    reg: f64,
    /// Source indices
    source_indices: Vec<usize>,
}

impl LcmvBeamformer {
    /// Create an LCMV beamformer from forward model and data covariance
    ///
    /// # Arguments
    /// * `gain` - Forward model gain matrix [n_channels x n_sources]
    /// * `data_cov` - Data covariance matrix
    /// * `noise_cov` - Optional noise covariance for regularization
    /// * `reg` - Regularization parameter (fraction of trace)
    pub fn make_lcmv(
        gain: &GainMatrix,
        data_cov: &Covariance,
        noise_cov: Option<&Covariance>,
        reg: f64,
    ) -> InverseResult<Self> {
        let n_channels = gain.n_sensors();
        let n_source_cols = gain.n_source_columns();

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

        // Get data covariance matrix
        let mut c = data_cov.data().clone();

        // Apply regularization
        if reg > 0.0 {
            let trace: f64 = (0..n_channels).map(|i| c[(i, i)]).sum();
            let reg_val = reg * trace / n_channels as f64;
            for i in 0..n_channels {
                c[(i, i)] += reg_val;
            }
        }

        // Add noise covariance if provided
        if let Some(nc) = noise_cov {
            if nc.n_channels() != n_channels {
                return Err(InverseError::DimensionMismatch(
                    "Noise covariance dimension mismatch".to_string(),
                ));
            }
            c += nc.data();
        }

        // Invert covariance matrix
        let c_inv = Self::invert_covariance(&c)?;

        // Convert gain to matrix
        let g = DMatrix::from_fn(n_channels, n_source_cols, |i, j| gain.data()[i][j]);

        // Compute filters: w_i = C^-1 * g_i / (g_i^T * C^-1 * g_i)
        let mut filters = DMatrix::zeros(n_source_cols, n_channels);

        for src in 0..n_source_cols {
            // Extract lead field for this source
            let g_i = g.column(src);

            // C^-1 * g
            let c_inv_g = &c_inv * g_i;

            // g^T * C^-1 * g (normalization factor)
            let norm = g_i.dot(&c_inv_g);

            if norm > 1e-30 {
                // w = C^-1 * g / norm
                for ch in 0..n_channels {
                    filters[(src, ch)] = c_inv_g[ch] / norm;
                }
            }
        }

        // For free orientation, n_sources = n_source_cols / 3
        let n_sources = if gain.is_free_orientation() {
            n_source_cols / 3
        } else {
            n_source_cols
        };

        Ok(Self {
            filters,
            n_sources,
            n_channels,
            reg,
            source_indices: (0..n_sources).collect(),
        })
    }

    /// Apply the beamformer to sensor data
    ///
    /// # Arguments
    /// * `data` - Sensor data [n_channels x n_times]
    ///
    /// # Returns
    /// Source estimates [n_sources x n_times]
    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();
        let n_source_cols = self.filters.nrows();

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

        // Apply filters: S = W * M
        let source_mat = &self.filters * &data_mat;

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

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

        let free_orientation = n_source_cols != self.n_sources;

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

    /// Get the spatial filters
    pub fn filters(&self) -> &DMatrix<f64> {
        &self.filters
    }

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

    /// Compute beamformer output power (pseudo-Z statistic)
    ///
    /// Returns the neural activity index (NAI) for each source
    pub fn compute_nai(
        &self,
        data_cov: &Covariance,
        noise_cov: &Covariance,
    ) -> InverseResult<Vec<f64>> {
        let n_source_cols = self.filters.nrows();
        let mut nai = Vec::with_capacity(n_source_cols);

        let c_data = data_cov.data();
        let c_noise = noise_cov.data();

        for src in 0..n_source_cols {
            // Extract filter row
            let w = self.filters.row(src);

            // Signal power: w^T * C_data * w
            let signal_power: f64 = (0..self.n_channels)
                .map(|i| {
                    (0..self.n_channels)
                        .map(|j| w[i] * c_data[(i, j)] * w[j])
                        .sum::<f64>()
                })
                .sum();

            // Noise power: w^T * C_noise * w
            let noise_power: f64 = (0..self.n_channels)
                .map(|i| {
                    (0..self.n_channels)
                        .map(|j| w[i] * c_noise[(i, j)] * w[j])
                        .sum::<f64>()
                })
                .sum();

            // NAI = signal / noise
            if noise_power > 1e-30 {
                nai.push(signal_power / noise_power);
            } else {
                nai.push(0.0);
            }
        }

        Ok(nai)
    }

    /// Invert covariance matrix using SVD
    fn invert_covariance(c: &DMatrix<f64>) -> InverseResult<DMatrix<f64>> {
        let n = c.nrows();

        let svd = c.clone().svd(true, true);
        let u = svd
            .u
            .ok_or_else(|| InverseError::ComputationError("SVD failed".to_string()))?;
        let vt = svd
            .v_t
            .ok_or_else(|| InverseError::ComputationError("SVD failed".to_string()))?;
        let s = svd.singular_values;

        let tol = 1e-10 * s[0];

        let s_inv = DMatrix::from_diagonal(&DVector::from_fn(n, |i, _| {
            if s[i] > tol { 1.0 / s[i] } else { 0.0 }
        }));

        Ok(vt.transpose() * &s_inv * u.transpose())
    }
}

// ============================================================================
// DICS Beamformer (Dynamic Imaging of Coherent Sources)
// ============================================================================

/// Configuration for DICS beamformer
#[derive(Debug, Clone)]
pub struct DicsConfig {
    /// Target frequencies (Hz)
    pub frequencies: Vec<f64>,
    /// Frequency bandwidth for smoothing (Hz)
    pub bandwidth: f64,
    /// Regularization parameter (fraction of trace)
    pub reg: f64,
    /// Pick orientation method
    pub pick_ori: PickOrientation,
    /// Whether to compute power (true) or complex values (false)
    pub real_filter: bool,
    /// Sampling frequency
    pub sfreq: f64,
}

impl Default for DicsConfig {
    fn default() -> Self {
        Self {
            frequencies: vec![10.0], // Default: alpha band
            bandwidth: 4.0,
            reg: 0.05,
            pick_ori: PickOrientation::MaxPower,
            real_filter: true,
            sfreq: 1000.0,
        }
    }
}

impl DicsConfig {
    /// Create DICS config for alpha band (8-12 Hz)
    pub fn alpha(sfreq: f64) -> Self {
        Self {
            frequencies: vec![10.0],
            bandwidth: 4.0,
            sfreq,
            ..Self::default()
        }
    }

    /// Create DICS config for beta band (13-30 Hz)
    pub fn beta(sfreq: f64) -> Self {
        Self {
            frequencies: vec![20.0],
            bandwidth: 17.0,
            sfreq,
            ..Self::default()
        }
    }

    /// Create DICS config for gamma band (30-100 Hz)
    pub fn gamma(sfreq: f64) -> Self {
        Self {
            frequencies: vec![50.0],
            bandwidth: 40.0,
            sfreq,
            ..Self::default()
        }
    }

    /// Set target frequencies
    pub fn with_frequencies(mut self, freqs: Vec<f64>) -> Self {
        self.frequencies = freqs;
        self
    }

    /// Set bandwidth
    pub fn with_bandwidth(mut self, bw: f64) -> Self {
        self.bandwidth = bw;
        self
    }

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

/// Orientation picking method for free-orientation sources
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PickOrientation {
    /// Pick orientation with maximum power
    MaxPower,
    /// Use surface normal (requires fixed orientation in gain)
    Normal,
    /// Return all orientations (no projection)
    None,
}

/// Cross-spectral density matrix
#[derive(Debug, Clone)]
pub struct CrossSpectralDensity {
    /// CSD matrix [n_channels x n_channels] (complex)
    data: Vec<Vec<Complex<f64>>>,
    /// Number of channels
    n_channels: usize,
    /// Frequencies at which CSD was computed
    frequencies: Vec<f64>,
    /// Number of epochs/segments used
    n_epochs: usize,
}

impl CrossSpectralDensity {
    /// Create a new CSD matrix
    pub fn new(data: Vec<Vec<Complex<f64>>>, frequencies: Vec<f64>, n_epochs: usize) -> Self {
        let n_channels = data.len();
        Self {
            data,
            n_channels,
            frequencies,
            n_epochs,
        }
    }

    /// Compute CSD from epochs of data
    ///
    /// # Arguments
    /// * `epochs` - Data epochs [n_epochs][n_channels][n_samples]
    /// * `sfreq` - Sampling frequency
    /// * `fmin` - Minimum frequency
    /// * `fmax` - Maximum frequency
    /// * `n_fft` - FFT size (None = use epoch length)
    pub fn from_epochs(
        epochs: &[Vec<Vec<f64>>],
        sfreq: f64,
        fmin: f64,
        fmax: f64,
        n_fft: Option<usize>,
    ) -> InverseResult<Self> {
        if epochs.is_empty() {
            return Err(InverseError::InvalidParameter(
                "No epochs provided".to_string(),
            ));
        }

        let n_epochs = epochs.len();
        let n_channels = epochs[0].len();
        let n_samples = epochs[0][0].len();
        let n_fft = n_fft.unwrap_or(n_samples);

        // Compute frequency bins
        let df = sfreq / n_fft as f64;
        let freq_min_idx = (fmin / df).ceil() as usize;
        let freq_max_idx = (fmax / df).floor() as usize;

        let frequencies: Vec<f64> = (freq_min_idx..=freq_max_idx)
            .map(|i| i as f64 * df)
            .collect();

        if frequencies.is_empty() {
            return Err(InverseError::InvalidParameter(
                "No frequencies in specified range".to_string(),
            ));
        }

        // Initialize CSD matrix
        let mut csd: Vec<Vec<Complex<f64>>> =
            vec![vec![Complex::new(0.0, 0.0); n_channels]; n_channels];

        // Compute FFT for each epoch and accumulate CSD
        for epoch in epochs {
            // Compute FFT for each channel
            let ffts: Vec<Vec<Complex<f64>>> = epoch
                .iter()
                .map(|ch| Self::compute_fft(ch, n_fft))
                .collect();

            // Accumulate cross-spectra for frequency range
            for i in 0..n_channels {
                for j in i..n_channels {
                    let mut sum = Complex::new(0.0, 0.0);
                    for &freq_idx in &(freq_min_idx..=freq_max_idx).collect::<Vec<_>>() {
                        if freq_idx < ffts[0].len() {
                            // CSD_ij = X_i * conj(X_j)
                            sum += ffts[i][freq_idx] * ffts[j][freq_idx].conj();
                        }
                    }
                    csd[i][j] += sum;
                    if i != j {
                        csd[j][i] += sum.conj();
                    }
                }
            }
        }

        // Normalize by number of epochs and frequency bins
        let n_freqs = frequencies.len() as f64;
        let norm = n_epochs as f64 * n_freqs;
        for i in 0..n_channels {
            for j in 0..n_channels {
                csd[i][j] /= norm;
            }
        }

        Ok(Self::new(csd, frequencies, n_epochs))
    }

    /// Get CSD data
    pub fn data(&self) -> &Vec<Vec<Complex<f64>>> {
        &self.data
    }

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

    /// Get frequencies
    pub fn frequencies(&self) -> &[f64] {
        &self.frequencies
    }

    /// Convert to real matrix (magnitude)
    pub fn to_real(&self) -> DMatrix<f64> {
        DMatrix::from_fn(self.n_channels, self.n_channels, |i, j| {
            self.data[i][j].norm()
        })
    }

    /// Convert to real matrix (real part only)
    pub fn real_part(&self) -> DMatrix<f64> {
        DMatrix::from_fn(self.n_channels, self.n_channels, |i, j| self.data[i][j].re)
    }

    /// Simple DFT implementation
    fn compute_fft(signal: &[f64], n_fft: usize) -> Vec<Complex<f64>> {
        let n = signal.len().min(n_fft);
        let mut result = vec![Complex::new(0.0, 0.0); n_fft / 2 + 1];

        for (k, res) in result.iter_mut().enumerate() {
            let mut sum = Complex::new(0.0, 0.0);
            for (t, &x) in signal.iter().take(n).enumerate() {
                let angle = -2.0 * PI * k as f64 * t as f64 / n_fft as f64;
                sum += Complex::new(x * angle.cos(), x * angle.sin());
            }
            *res = sum / (n as f64).sqrt();
        }

        result
    }
}

/// DICS beamformer
#[derive(Debug, Clone)]
pub struct DicsBeamformer {
    /// Spatial filters [n_sources x n_channels] (complex)
    filters: Vec<Vec<Complex<f64>>>,
    /// Number of sources
    n_sources: usize,
    /// Number of channels
    n_channels: usize,
    /// Whether sources have free orientation
    free_orientation: bool,
    /// Configuration used
    config: DicsConfig,
    /// Source indices
    source_indices: Vec<usize>,
}

impl DicsBeamformer {
    /// Create a DICS beamformer from forward model and CSD
    ///
    /// # Arguments
    /// * `gain` - Forward model gain matrix
    /// * `csd` - Cross-spectral density matrix
    /// * `noise_csd` - Optional noise CSD for regularization
    /// * `config` - DICS configuration
    pub fn make_dics(
        gain: &GainMatrix,
        csd: &CrossSpectralDensity,
        noise_csd: Option<&CrossSpectralDensity>,
        config: DicsConfig,
    ) -> InverseResult<Self> {
        let n_channels = gain.n_sensors();
        let n_source_cols = gain.n_source_columns();
        let free_orientation = gain.is_free_orientation();

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

        // Convert CSD to matrix form (use real part for real filter)
        let csd_mat = if config.real_filter {
            csd.real_part()
        } else {
            csd.to_real()
        };

        // Apply regularization
        let mut c = csd_mat.clone();
        if config.reg > 0.0 {
            let trace: f64 = (0..n_channels).map(|i| c[(i, i)]).sum();
            let reg_val = config.reg * trace / n_channels as f64;
            for i in 0..n_channels {
                c[(i, i)] += reg_val;
            }
        }

        // Add noise CSD if provided
        if let Some(nc) = noise_csd {
            let noise_mat = if config.real_filter {
                nc.real_part()
            } else {
                nc.to_real()
            };
            c += &noise_mat;
        }

        // Invert CSD matrix
        let c_inv = Self::invert_matrix(&c)?;

        // Convert gain to matrix
        let g = DMatrix::from_fn(n_channels, n_source_cols, |i, j| gain.data()[i][j]);

        // Compute filters
        let filters = if free_orientation {
            Self::compute_filters_free(&g, &c_inv, n_channels, n_source_cols, &config)?
        } else {
            Self::compute_filters_fixed(&g, &c_inv, n_channels, n_source_cols)?
        };

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

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

    /// Compute source power from CSD
    ///
    /// Returns power at each source location
    pub fn compute_power(&self, csd: &CrossSpectralDensity) -> InverseResult<Vec<f64>> {
        if csd.n_channels() != self.n_channels {
            return Err(InverseError::DimensionMismatch(
                "CSD channel mismatch".to_string(),
            ));
        }

        let n_source_cols = self.filters.len();
        let mut power = Vec::with_capacity(n_source_cols);

        for src in 0..n_source_cols {
            let w = &self.filters[src];

            // Power = w^H * CSD * w
            let mut p = Complex::new(0.0, 0.0);
            for i in 0..self.n_channels {
                for j in 0..self.n_channels {
                    p += w[i].conj() * csd.data()[i][j] * w[j];
                }
            }

            power.push(p.re.max(0.0)); // Power is real and non-negative
        }

        // For free orientation, combine orientations
        if self.free_orientation && n_source_cols > self.n_sources {
            let mut combined = Vec::with_capacity(self.n_sources);
            for src in 0..self.n_sources {
                let p = power[3 * src] + power[3 * src + 1] + power[3 * src + 2];
                combined.push(p);
            }
            return Ok(combined);
        }

        Ok(power)
    }

    /// Apply beamformer to time-domain data
    ///
    /// First computes CSD from data, then returns source power
    pub fn apply_to_epochs(&self, epochs: &[Vec<Vec<f64>>]) -> InverseResult<Vec<f64>> {
        // Compute CSD from epochs
        let fmin =
            self.config.frequencies.first().copied().unwrap_or(1.0) - self.config.bandwidth / 2.0;
        let fmax =
            self.config.frequencies.last().copied().unwrap_or(100.0) + self.config.bandwidth / 2.0;

        let csd = CrossSpectralDensity::from_epochs(
            epochs,
            self.config.sfreq,
            fmin.max(0.1),
            fmax,
            None,
        )?;

        self.compute_power(&csd)
    }

    /// Get the spatial filters
    pub fn filters(&self) -> &Vec<Vec<Complex<f64>>> {
        &self.filters
    }

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

    /// Get frequencies
    pub fn frequencies(&self) -> &[f64] {
        &self.config.frequencies
    }

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

    /// Compute filters for fixed orientation
    fn compute_filters_fixed(
        g: &DMatrix<f64>,
        c_inv: &DMatrix<f64>,
        n_channels: usize,
        n_sources: usize,
    ) -> InverseResult<Vec<Vec<Complex<f64>>>> {
        let mut filters = Vec::with_capacity(n_sources);

        for src in 0..n_sources {
            let g_i = g.column(src);
            let c_inv_g = c_inv * g_i;
            let norm = g_i.dot(&c_inv_g);

            let filter: Vec<Complex<f64>> = if norm > 1e-30 {
                (0..n_channels)
                    .map(|ch| Complex::new(c_inv_g[ch] / norm, 0.0))
                    .collect()
            } else {
                vec![Complex::new(0.0, 0.0); n_channels]
            };

            filters.push(filter);
        }

        Ok(filters)
    }

    /// Compute filters for free orientation
    fn compute_filters_free(
        g: &DMatrix<f64>,
        c_inv: &DMatrix<f64>,
        n_channels: usize,
        n_source_cols: usize,
        config: &DicsConfig,
    ) -> InverseResult<Vec<Vec<Complex<f64>>>> {
        let n_sources = n_source_cols / 3;
        let mut filters = Vec::with_capacity(n_source_cols);

        match config.pick_ori {
            PickOrientation::MaxPower => {
                // Compute filter for each orientation, then pick max power
                for src in 0..n_sources {
                    // Compute 3x3 matrix: G^T * C^(-1) * G
                    let mut gtcg = [[0.0; 3]; 3];
                    for ori1 in 0..3 {
                        let col1 = 3 * src + ori1;
                        for ori2 in 0..3 {
                            let col2 = 3 * src + ori2;
                            for i in 0..n_channels {
                                for j in 0..n_channels {
                                    gtcg[ori1][ori2] += g[(i, col1)] * c_inv[(i, j)] * g[(j, col2)];
                                }
                            }
                        }
                    }

                    // Find eigenvector with largest eigenvalue
                    let (max_ori, _) = Self::max_eigenvector_3x3(&gtcg);

                    // Create combined leadfield
                    let g_combined: Vec<f64> = (0..n_channels)
                        .map(|ch| {
                            max_ori[0] * g[(ch, 3 * src)]
                                + max_ori[1] * g[(ch, 3 * src + 1)]
                                + max_ori[2] * g[(ch, 3 * src + 2)]
                        })
                        .collect();

                    // Compute filter for combined orientation
                    let g_vec = DVector::from_column_slice(&g_combined);
                    let c_inv_g = c_inv * &g_vec;
                    let norm = g_vec.dot(&c_inv_g);

                    let filter: Vec<Complex<f64>> = if norm > 1e-30 {
                        (0..n_channels)
                            .map(|ch| Complex::new(c_inv_g[ch] / norm, 0.0))
                            .collect()
                    } else {
                        vec![Complex::new(0.0, 0.0); n_channels]
                    };

                    // Store same filter for all orientations (max power direction)
                    filters.push(filter.clone());
                    filters.push(filter.clone());
                    filters.push(filter);
                }
            }
            PickOrientation::None | PickOrientation::Normal => {
                // Compute separate filter for each orientation
                for src_col in 0..n_source_cols {
                    let g_i = g.column(src_col);
                    let c_inv_g = c_inv * g_i;
                    let norm = g_i.dot(&c_inv_g);

                    let filter: Vec<Complex<f64>> = if norm > 1e-30 {
                        (0..n_channels)
                            .map(|ch| Complex::new(c_inv_g[ch] / norm, 0.0))
                            .collect()
                    } else {
                        vec![Complex::new(0.0, 0.0); n_channels]
                    };

                    filters.push(filter);
                }
            }
        }

        Ok(filters)
    }

    /// Find eigenvector with maximum eigenvalue for 3x3 symmetric matrix
    fn max_eigenvector_3x3(a: &[[f64; 3]; 3]) -> ([f64; 3], f64) {
        // Power iteration for dominant eigenvector
        let mut v = [1.0 / 3.0_f64.sqrt(); 3];

        for _ in 0..20 {
            // Av
            let mut av = [0.0; 3];
            for i in 0..3 {
                for j in 0..3 {
                    av[i] += a[i][j] * v[j];
                }
            }

            // Normalize
            let norm = (av[0] * av[0] + av[1] * av[1] + av[2] * av[2]).sqrt();
            if norm > 1e-15 {
                for i in 0..3 {
                    v[i] = av[i] / norm;
                }
            }
        }

        // Compute eigenvalue
        let mut av = [0.0; 3];
        for i in 0..3 {
            for j in 0..3 {
                av[i] += a[i][j] * v[j];
            }
        }
        let eigenvalue = v[0] * av[0] + v[1] * av[1] + v[2] * av[2];

        (v, eigenvalue)
    }

    /// Invert matrix using SVD
    fn invert_matrix(c: &DMatrix<f64>) -> InverseResult<DMatrix<f64>> {
        let n = c.nrows();

        let svd = c.clone().svd(true, true);
        let u = svd
            .u
            .ok_or_else(|| InverseError::ComputationError("SVD failed".to_string()))?;
        let vt = svd
            .v_t
            .ok_or_else(|| InverseError::ComputationError("SVD failed".to_string()))?;
        let s = svd.singular_values;

        let tol = 1e-10 * s[0];

        let s_inv = DMatrix::from_diagonal(&DVector::from_fn(n, |i, _| {
            if s[i] > tol { 1.0 / s[i] } else { 0.0 }
        }));

        Ok(vt.transpose() * &s_inv * u.transpose())
    }
}

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

    fn create_simple_gain() -> GainMatrix {
        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()
    }

    fn create_free_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()
    }

    #[test]
    fn test_make_lcmv() {
        let gain = create_simple_gain();
        let data_cov = Covariance::identity(4, CovarianceType::Data);

        let bf = LcmvBeamformer::make_lcmv(&gain, &data_cov, None, 0.05).unwrap();

        assert_eq!(bf.n_sources(), 3);
    }

    #[test]
    fn test_apply_lcmv() {
        let gain = create_simple_gain();
        let data_cov = Covariance::identity(4, CovarianceType::Data);

        let bf = LcmvBeamformer::make_lcmv(&gain, &data_cov, None, 0.05).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 = bf.apply(&data).unwrap();
        assert_eq!(stc.n_sources(), 3);
        assert_eq!(stc.n_times(), 2);
    }

    // DICS tests

    #[test]
    fn test_dics_config() {
        let config = DicsConfig::default();
        assert_eq!(config.frequencies, vec![10.0]);
        assert!((config.bandwidth - 4.0).abs() < 1e-10);

        let config = DicsConfig::alpha(1000.0).with_reg(0.1);
        assert!((config.reg - 0.1).abs() < 1e-10);
        assert!((config.sfreq - 1000.0).abs() < 1e-10);
    }

    #[test]
    fn test_csd_from_epochs() {
        // Create simple test epochs
        let n_epochs = 5;
        let n_channels = 4;
        let n_samples = 100;

        let epochs: Vec<Vec<Vec<f64>>> = (0..n_epochs)
            .map(|_| {
                (0..n_channels)
                    .map(|ch| {
                        (0..n_samples)
                            .map(|t| {
                                // Simple sinusoid + noise
                                let freq = 10.0;
                                let sfreq = 100.0;
                                (2.0 * PI * freq * t as f64 / sfreq).sin()
                                    + 0.1 * (ch as f64 * t as f64 * 0.01).sin()
                            })
                            .collect()
                    })
                    .collect()
            })
            .collect();

        let csd = CrossSpectralDensity::from_epochs(
            &epochs, 100.0, // sfreq
            5.0,   // fmin
            15.0,  // fmax
            None,
        )
        .unwrap();

        assert_eq!(csd.n_channels(), 4);
        assert!(!csd.frequencies().is_empty());
    }

    #[test]
    fn test_make_dics_fixed() {
        let gain = create_simple_gain();

        // Create simple CSD (identity-like for testing)
        let n_channels = 4;
        let csd_data: Vec<Vec<Complex<f64>>> = (0..n_channels)
            .map(|i| {
                (0..n_channels)
                    .map(|j| {
                        if i == j {
                            Complex::new(1.0, 0.0)
                        } else {
                            Complex::new(0.1, 0.0)
                        }
                    })
                    .collect()
            })
            .collect();

        let csd = CrossSpectralDensity::new(csd_data, vec![10.0], 10);

        let config = DicsConfig::default();
        let bf = DicsBeamformer::make_dics(&gain, &csd, None, config).unwrap();

        assert_eq!(bf.n_sources(), 3);
    }

    #[test]
    fn test_make_dics_free() {
        let gain = create_free_gain();

        let n_channels = 4;
        let csd_data: Vec<Vec<Complex<f64>>> = (0..n_channels)
            .map(|i| {
                (0..n_channels)
                    .map(|j| {
                        if i == j {
                            Complex::new(1.0, 0.0)
                        } else {
                            Complex::new(0.1, 0.0)
                        }
                    })
                    .collect()
            })
            .collect();

        let csd = CrossSpectralDensity::new(csd_data, vec![10.0], 10);

        let config = DicsConfig::default();
        let bf = DicsBeamformer::make_dics(&gain, &csd, None, config).unwrap();

        assert_eq!(bf.n_sources(), 2); // 6 columns / 3 orientations
    }

    #[test]
    fn test_dics_power() {
        let gain = create_simple_gain();

        let n_channels = 4;
        let csd_data: Vec<Vec<Complex<f64>>> = (0..n_channels)
            .map(|i| {
                (0..n_channels)
                    .map(|j| {
                        if i == j {
                            Complex::new(1.0, 0.0)
                        } else {
                            Complex::new(0.1, 0.01)
                        }
                    })
                    .collect()
            })
            .collect();

        let csd = CrossSpectralDensity::new(csd_data, vec![10.0], 10);

        let config = DicsConfig::default();
        let bf = DicsBeamformer::make_dics(&gain, &csd, None, config).unwrap();

        let power = bf.compute_power(&csd).unwrap();
        assert_eq!(power.len(), 3);

        // All power values should be non-negative
        for p in &power {
            assert!(*p >= 0.0);
        }
    }
}