//! Noise and data covariance estimation.
//!
//! Covariance matrices are essential for inverse solutions:
//! - **Noise covariance**: Estimated from baseline or empty room recordings
//! - **Data covariance**: Used for beamforming approaches

use crate::{InverseError, InverseResult};
use nalgebra::{DMatrix, DVector};

/// Type of covariance matrix
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CovarianceType {
    /// Noise covariance (from baseline)
    Noise,
    /// Data covariance (from epochs)
    Data,
    /// Empty room noise
    EmptyRoom,
}

/// Covariance matrix with metadata
#[derive(Debug, Clone)]
pub struct Covariance {
    /// Covariance matrix [n_channels x n_channels]
    data: DMatrix<f64>,
    /// Channel names
    channel_names: Vec<String>,
    /// Type of covariance
    cov_type: CovarianceType,
    /// Number of samples used
    n_samples: usize,
    /// Regularization applied
    reg: f64,
}

impl Covariance {
    /// Create a covariance from a matrix
    pub fn new(
        data: Vec<Vec<f64>>,
        channel_names: Vec<String>,
        cov_type: CovarianceType,
        n_samples: usize,
    ) -> InverseResult<Self> {
        let n = data.len();
        if n == 0 {
            return Err(InverseError::InvalidParameter(
                "Covariance matrix cannot be empty".to_string(),
            ));
        }

        if data[0].len() != n {
            return Err(InverseError::DimensionMismatch(
                "Covariance matrix must be square".to_string(),
            ));
        }

        if channel_names.len() != n {
            return Err(InverseError::DimensionMismatch(format!(
                "Expected {} channel names, got {}",
                n,
                channel_names.len()
            )));
        }

        // Convert to nalgebra matrix
        let mat = DMatrix::from_fn(n, n, |i, j| data[i][j]);

        Ok(Self {
            data: mat,
            channel_names,
            cov_type,
            n_samples,
            reg: 0.0,
        })
    }

    /// Create an identity covariance (for testing)
    pub fn identity(n_channels: usize, cov_type: CovarianceType) -> Self {
        let names: Vec<String> = (0..n_channels).map(|i| format!("CH{:03}", i + 1)).collect();

        Self {
            data: DMatrix::identity(n_channels, n_channels),
            channel_names: names,
            cov_type,
            n_samples: 0,
            reg: 0.0,
        }
    }

    /// Compute covariance from data epochs
    ///
    /// # Arguments
    /// * `epochs` - Data epochs [n_epochs][n_channels][n_times]
    /// * `channel_names` - Channel names
    /// * `cov_type` - Type of covariance
    /// * `tmin` - Start time index for baseline
    /// * `tmax` - End time index for baseline
    pub fn from_epochs(
        epochs: &[Vec<Vec<f64>>],
        channel_names: Vec<String>,
        cov_type: CovarianceType,
        tmin: Option<usize>,
        tmax: Option<usize>,
    ) -> InverseResult<Self> {
        if epochs.is_empty() {
            return Err(InverseError::InvalidParameter(
                "No epochs provided".to_string(),
            ));
        }

        let n_channels = epochs[0].len();
        let n_times_total = epochs[0][0].len();

        // Determine time range
        let t_start = tmin.unwrap_or(0);
        let t_end = tmax.unwrap_or(n_times_total);

        // Collect all samples across epochs and time
        let mut all_samples: Vec<Vec<f64>> = Vec::new();

        for epoch in epochs {
            for t in t_start..t_end {
                let sample: Vec<f64> = epoch.iter().map(|ch| ch[t]).collect();
                all_samples.push(sample);
            }
        }

        let n_samples = all_samples.len();

        // Compute mean
        let mut mean = vec![0.0; n_channels];
        for sample in &all_samples {
            for (i, &val) in sample.iter().enumerate() {
                mean[i] += val;
            }
        }
        for m in &mut mean {
            *m /= n_samples as f64;
        }

        // Compute covariance
        let mut cov = vec![vec![0.0; n_channels]; n_channels];
        for sample in &all_samples {
            for i in 0..n_channels {
                for j in i..n_channels {
                    cov[i][j] += (sample[i] - mean[i]) * (sample[j] - mean[j]);
                }
            }
        }

        // Normalize and symmetrize
        for i in 0..n_channels {
            for j in i..n_channels {
                cov[i][j] /= (n_samples - 1) as f64;
                cov[j][i] = cov[i][j];
            }
        }

        Self::new(cov, channel_names, cov_type, n_samples)
    }

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

    /// Get channel names
    pub fn channel_names(&self) -> &[String] {
        &self.channel_names
    }

    /// Get covariance type
    pub fn cov_type(&self) -> CovarianceType {
        self.cov_type
    }

    /// Get number of samples used
    pub fn n_samples(&self) -> usize {
        self.n_samples
    }

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

    /// Get a copy of the data as nested vectors
    pub fn to_vec(&self) -> Vec<Vec<f64>> {
        let n = self.n_channels();
        (0..n)
            .map(|i| (0..n).map(|j| self.data[(i, j)]).collect())
            .collect()
    }

    /// Apply Tikhonov regularization
    ///
    /// C_reg = C + reg * trace(C)/n * I
    pub fn regularize(&mut self, reg: f64) -> InverseResult<()> {
        if !(0.0..=1.0).contains(&reg) {
            return Err(InverseError::InvalidParameter(
                "Regularization must be between 0 and 1".to_string(),
            ));
        }

        let n = self.n_channels();
        let trace: f64 = (0..n).map(|i| self.data[(i, i)]).sum();
        let reg_value = reg * trace / n as f64;

        for i in 0..n {
            self.data[(i, i)] += reg_value;
        }

        self.reg = reg;
        Ok(())
    }

    /// Compute the whitening matrix W = C^(-1/2)
    ///
    /// Such that W * C * W^T = I
    pub fn compute_whitener(&self) -> InverseResult<DMatrix<f64>> {
        let n = self.n_channels();

        // Eigendecomposition: C = V * D * V^T
        let eigen = self.data.clone().symmetric_eigen();
        let eigenvalues = eigen.eigenvalues;
        let eigenvectors = eigen.eigenvectors;

        // Check for negative eigenvalues
        let min_eig = eigenvalues.min();
        if min_eig <= 0.0 {
            return Err(InverseError::ComputationError(format!(
                "Covariance has non-positive eigenvalue: {:.2e}. Consider regularization.",
                min_eig
            )));
        }

        // W = V * D^(-1/2) * V^T
        let d_inv_sqrt =
            DMatrix::from_diagonal(&DVector::from_fn(n, |i, _| 1.0 / eigenvalues[i].sqrt()));

        let whitener = &eigenvectors * &d_inv_sqrt * eigenvectors.transpose();
        Ok(whitener)
    }

    /// Compute the pseudo-inverse using SVD with rank control
    pub fn compute_pseudo_inverse(&self, rank: Option<usize>) -> InverseResult<DMatrix<f64>> {
        let n = self.n_channels();

        let svd = self.data.clone().svd(true, true);
        let u = svd
            .u
            .ok_or_else(|| InverseError::ComputationError("SVD failed to compute U".to_string()))?;
        let vt = svd.v_t.ok_or_else(|| {
            InverseError::ComputationError("SVD failed to compute V^T".to_string())
        })?;
        let s = svd.singular_values;

        // Determine rank
        let effective_rank = rank.unwrap_or_else(|| {
            let tol = 1e-10 * s[0];
            s.iter().filter(|&&v| v > tol).count()
        });

        // Compute pseudo-inverse: V * S^(-1) * U^T (only for non-zero singular values)
        let s_inv = DMatrix::from_diagonal(&DVector::from_fn(n, |i, _| {
            if i < effective_rank && s[i] > 1e-15 {
                1.0 / s[i]
            } else {
                0.0
            }
        }));

        let pinv = vt.transpose() * &s_inv * u.transpose();
        Ok(pinv)
    }
}

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

    #[test]
    fn test_identity_covariance() {
        let cov = Covariance::identity(10, CovarianceType::Noise);
        assert_eq!(cov.n_channels(), 10);
        assert_eq!(cov.cov_type(), CovarianceType::Noise);
    }

    #[test]
    fn test_from_epochs() {
        // Simple test data: 2 epochs, 3 channels, 5 time points
        let epochs = vec![
            vec![
                vec![1.0, 2.0, 3.0, 4.0, 5.0],
                vec![2.0, 3.0, 4.0, 5.0, 6.0],
                vec![3.0, 4.0, 5.0, 6.0, 7.0],
            ],
            vec![
                vec![1.5, 2.5, 3.5, 4.5, 5.5],
                vec![2.5, 3.5, 4.5, 5.5, 6.5],
                vec![3.5, 4.5, 5.5, 6.5, 7.5],
            ],
        ];
        let names = vec!["CH1".to_string(), "CH2".to_string(), "CH3".to_string()];

        let cov =
            Covariance::from_epochs(&epochs, names, CovarianceType::Data, None, None).unwrap();

        assert_eq!(cov.n_channels(), 3);
        assert_eq!(cov.n_samples(), 10); // 2 epochs * 5 time points
    }

    #[test]
    fn test_regularization() {
        let mut cov = Covariance::identity(5, CovarianceType::Noise);
        cov.regularize(0.1).unwrap();

        // Diagonal should be > 1 after regularization
        let data = cov.data();
        assert!(data[(0, 0)] > 1.0);
    }

    #[test]
    fn test_whitener() {
        // 2x2 covariance matrix
        let data = vec![vec![4.0, 2.0], vec![2.0, 3.0]];
        let names = vec!["A".to_string(), "B".to_string()];
        let cov = Covariance::new(data, names, CovarianceType::Noise, 100).unwrap();

        let whitener = cov.compute_whitener().unwrap();

        // W * C * W^T should be close to identity
        let result = &whitener * cov.data() * whitener.transpose();
        assert!((result[(0, 0)] - 1.0).abs() < 1e-10);
        assert!((result[(1, 1)] - 1.0).abs() < 1e-10);
        assert!((result[(0, 1)]).abs() < 1e-10);
    }
}