//! Independent Component Analysis (ICA) for artifact removal.
//!
//! ICA separates the signal into statistically independent components,
//! which can be used to identify and remove artifacts like eye blinks,
//! muscle activity, and heartbeat.
//!
//! ## Mathematical Background
//!
//! The mixing model is: X = A * S
//! where X is the observed data, A is the mixing matrix, and S are sources.
//!
//! FastICA finds the unmixing matrix W such that: S = W * X
//! by maximizing non-Gaussianity (negentropy).

use crate::{SignalError, SignalResult};

/// ICA estimation method
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum IcaMethod {
    /// FastICA with logcosh nonlinearity
    #[default]
    FastIcaLogcosh,
    /// FastICA with exponential nonlinearity
    FastIcaExp,
    /// FastICA with cubic nonlinearity
    FastIcaCube,
}

/// ICA decomposition result
#[derive(Debug, Clone)]
pub struct Ica {
    /// Mixing matrix A [n_channels x n_components]
    mixing: Vec<Vec<f64>>,
    /// Unmixing matrix W [n_components x n_channels]
    unmixing: Vec<Vec<f64>>,
    /// Whitening matrix
    whitening: Vec<Vec<f64>>,
    /// Mean of the data (for centering)
    mean: Vec<f64>,
    /// Number of components
    n_components: usize,
    /// Number of channels
    n_channels: usize,
    /// Method used
    method: IcaMethod,
    /// Explained variance ratio per component
    explained_variance_ratio: Vec<f64>,
}

impl Ica {
    /// Fit ICA to data
    ///
    /// # Arguments
    /// * `data` - Data matrix [n_channels][n_samples]
    /// * `n_components` - Number of components (None = n_channels)
    /// * `method` - ICA method
    /// * `max_iter` - Maximum iterations
    /// * `tol` - Convergence tolerance
    pub fn fit(
        data: &[Vec<f64>],
        n_components: Option<usize>,
        method: IcaMethod,
        max_iter: usize,
        tol: f64,
    ) -> SignalResult<Self> {
        if data.is_empty() || data[0].is_empty() {
            return Err(SignalError::InvalidLength("Empty data".to_string()));
        }

        let n_channels = data.len();
        let n_samples = data[0].len();
        let n_components = n_components.unwrap_or(n_channels).min(n_channels);

        // Center the data
        let mean: Vec<f64> = data
            .iter()
            .map(|ch| ch.iter().sum::<f64>() / n_samples as f64)
            .collect();

        let centered: Vec<Vec<f64>> = data
            .iter()
            .zip(&mean)
            .map(|(ch, &m)| ch.iter().map(|&x| x - m).collect())
            .collect();

        // Whiten the data (PCA + normalization)
        let (whitened, whitening, explained_var) = whiten(&centered, n_components)?;

        // FastICA
        let unmixing_white = fastica(&whitened, method, max_iter, tol)?;

        // Compute full unmixing matrix: W = W_ica * W_white
        let unmixing = mat_mult(&unmixing_white, &whitening);

        // Compute mixing matrix: A = W^(-1) = W_white^(-1) * W_ica^(-1)
        let mixing = pseudo_inverse(&unmixing)?;

        Ok(Self {
            mixing,
            unmixing,
            whitening,
            mean,
            n_components,
            n_channels,
            method,
            explained_variance_ratio: explained_var,
        })
    }

    /// Transform data to independent components
    ///
    /// # Arguments
    /// * `data` - Data matrix [n_channels][n_samples]
    ///
    /// # Returns
    /// Independent components [n_components][n_samples]
    pub fn transform(&self, data: &[Vec<f64>]) -> SignalResult<Vec<Vec<f64>>> {
        if data.len() != self.n_channels {
            return Err(SignalError::InvalidLength(format!(
                "Expected {} channels, got {}",
                self.n_channels,
                data.len()
            )));
        }

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

        // Center
        let centered: Vec<Vec<f64>> = data
            .iter()
            .zip(&self.mean)
            .map(|(ch, &m)| ch.iter().map(|&x| x - m).collect())
            .collect();

        // Apply unmixing: S = W * X
        let mut sources = vec![vec![0.0; n_samples]; self.n_components];
        for comp in 0..self.n_components {
            for t in 0..n_samples {
                for ch in 0..self.n_channels {
                    sources[comp][t] += self.unmixing[comp][ch] * centered[ch][t];
                }
            }
        }

        Ok(sources)
    }

    /// Inverse transform: reconstruct data from components
    ///
    /// # Arguments
    /// * `sources` - Independent components [n_components][n_samples]
    ///
    /// # Returns
    /// Reconstructed data [n_channels][n_samples]
    pub fn inverse_transform(&self, sources: &[Vec<f64>]) -> SignalResult<Vec<Vec<f64>>> {
        if sources.len() != self.n_components {
            return Err(SignalError::InvalidLength(format!(
                "Expected {} components, got {}",
                self.n_components,
                sources.len()
            )));
        }

        let n_samples = sources[0].len();

        // Apply mixing: X = A * S + mean
        let mut data = vec![vec![0.0; n_samples]; self.n_channels];
        for ch in 0..self.n_channels {
            for t in 0..n_samples {
                for comp in 0..self.n_components {
                    data[ch][t] += self.mixing[ch][comp] * sources[comp][t];
                }
                data[ch][t] += self.mean[ch];
            }
        }

        Ok(data)
    }

    /// Apply ICA to remove specific components
    ///
    /// # Arguments
    /// * `data` - Data matrix [n_channels][n_samples]
    /// * `exclude` - Indices of components to remove
    ///
    /// # Returns
    /// Cleaned data with specified components removed
    pub fn apply(&self, data: &[Vec<f64>], exclude: &[usize]) -> SignalResult<Vec<Vec<f64>>> {
        // Get sources
        let mut sources = self.transform(data)?;

        // Zero out excluded components
        for &idx in exclude {
            if idx < self.n_components {
                for x in &mut sources[idx] {
                    *x = 0.0;
                }
            }
        }

        // Reconstruct
        self.inverse_transform(&sources)
    }

    /// Get independent components from the data used for fitting
    pub fn get_sources(&self, data: &[Vec<f64>]) -> SignalResult<Vec<Vec<f64>>> {
        self.transform(data)
    }

    /// Get mixing matrix A
    pub fn mixing(&self) -> &[Vec<f64>] {
        &self.mixing
    }

    /// Get unmixing matrix W
    pub fn unmixing(&self) -> &[Vec<f64>] {
        &self.unmixing
    }

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

    /// Get explained variance ratio
    pub fn explained_variance_ratio(&self) -> &[f64] {
        &self.explained_variance_ratio
    }

    /// Compute component properties for artifact detection
    ///
    /// Returns kurtosis for each component (high kurtosis often indicates artifacts)
    pub fn component_kurtosis(&self, data: &[Vec<f64>]) -> SignalResult<Vec<f64>> {
        let sources = self.transform(data)?;

        let kurtosis: Vec<f64> = sources
            .iter()
            .map(|comp| {
                let n = comp.len() as f64;
                let mean = comp.iter().sum::<f64>() / n;
                let var = comp.iter().map(|&x| (x - mean).powi(2)).sum::<f64>() / n;
                if var < 1e-15 {
                    return 0.0;
                }
                let std = var.sqrt();
                let m4 = comp
                    .iter()
                    .map(|&x| ((x - mean) / std).powi(4))
                    .sum::<f64>()
                    / n;
                m4 - 3.0 // Excess kurtosis
            })
            .collect();

        Ok(kurtosis)
    }

    /// Find components correlated with a reference signal (e.g., EOG, ECG)
    ///
    /// Returns correlation coefficient for each component
    pub fn find_correlated_components(
        &self,
        data: &[Vec<f64>],
        reference: &[f64],
    ) -> SignalResult<Vec<f64>> {
        let sources = self.transform(data)?;

        if reference.len() != sources[0].len() {
            return Err(SignalError::InvalidLength(
                "Reference length must match data length".to_string(),
            ));
        }

        let n = reference.len() as f64;
        let ref_mean = reference.iter().sum::<f64>() / n;
        let ref_var: f64 = reference.iter().map(|&x| (x - ref_mean).powi(2)).sum();

        let correlations: Vec<f64> = sources
            .iter()
            .map(|comp| {
                let comp_mean = comp.iter().sum::<f64>() / n;
                let comp_var: f64 = comp.iter().map(|&x| (x - comp_mean).powi(2)).sum();

                if comp_var < 1e-15 || ref_var < 1e-15 {
                    return 0.0;
                }

                let cov: f64 = comp
                    .iter()
                    .zip(reference)
                    .map(|(&c, &r)| (c - comp_mean) * (r - ref_mean))
                    .sum();

                cov / (comp_var * ref_var).sqrt()
            })
            .collect();

        Ok(correlations)
    }
}

/// Whiten data using PCA
fn whiten(
    data: &[Vec<f64>],
    n_components: usize,
) -> SignalResult<(Vec<Vec<f64>>, Vec<Vec<f64>>, Vec<f64>)> {
    let n_channels = data.len();
    let n_samples = data[0].len();

    // Compute covariance matrix
    let mut cov = vec![vec![0.0; n_channels]; n_channels];
    for i in 0..n_channels {
        for j in i..n_channels {
            let c: f64 =
                (0..n_samples).map(|t| data[i][t] * data[j][t]).sum::<f64>() / n_samples as f64;
            cov[i][j] = c;
            cov[j][i] = c;
        }
    }

    // Eigendecomposition
    let (eigenvectors, eigenvalues) = symmetric_eigen(&cov)?;

    // Compute explained variance
    let total_var: f64 = eigenvalues.iter().sum();
    let explained_var: Vec<f64> = eigenvalues
        .iter()
        .take(n_components)
        .map(|&v| v / total_var)
        .collect();

    // Build whitening matrix: W = D^(-1/2) * V^T
    let mut whitening = vec![vec![0.0; n_channels]; n_components];
    for i in 0..n_components {
        let scale = if eigenvalues[i] > 1e-10 {
            1.0 / eigenvalues[i].sqrt()
        } else {
            0.0
        };
        for j in 0..n_channels {
            whitening[i][j] = scale * eigenvectors[i][j];
        }
    }

    // Apply whitening
    let mut whitened = vec![vec![0.0; n_samples]; n_components];
    for comp in 0..n_components {
        for t in 0..n_samples {
            for ch in 0..n_channels {
                whitened[comp][t] += whitening[comp][ch] * data[ch][t];
            }
        }
    }

    Ok((whitened, whitening, explained_var))
}

/// FastICA algorithm
fn fastica(
    whitened: &[Vec<f64>],
    method: IcaMethod,
    max_iter: usize,
    tol: f64,
) -> SignalResult<Vec<Vec<f64>>> {
    let n_components = whitened.len();
    let n_samples = whitened[0].len();

    // Initialize unmixing matrix with random-ish values
    let mut w = vec![vec![0.0; n_components]; n_components];
    for i in 0..n_components {
        for j in 0..n_components {
            w[i][j] = ((i * 7 + j * 13 + 1) as f64).sin();
        }
    }

    // Orthogonalize
    orthogonalize(&mut w);

    // FastICA iteration
    for _ in 0..max_iter {
        let w_old = w.clone();

        for p in 0..n_components {
            // Compute w^T * x for all samples
            let wx: Vec<f64> = (0..n_samples)
                .map(|t| (0..n_components).map(|i| w[p][i] * whitened[i][t]).sum())
                .collect();

            // Compute g(w^T * x) and g'(w^T * x)
            let (g, g_prime) = match method {
                IcaMethod::FastIcaLogcosh => {
                    let g: Vec<f64> = wx.iter().map(|&x| x.tanh()).collect();
                    let g_p: Vec<f64> = wx.iter().map(|&x| 1.0 - x.tanh().powi(2)).collect();
                    (g, g_p)
                }
                IcaMethod::FastIcaExp => {
                    let g: Vec<f64> = wx.iter().map(|&x| x * (-x * x / 2.0).exp()).collect();
                    let g_p: Vec<f64> = wx
                        .iter()
                        .map(|&x| (1.0 - x * x) * (-x * x / 2.0).exp())
                        .collect();
                    (g, g_p)
                }
                IcaMethod::FastIcaCube => {
                    let g: Vec<f64> = wx.iter().map(|&x| x * x * x).collect();
                    let g_p: Vec<f64> = wx.iter().map(|&x| 3.0 * x * x).collect();
                    (g, g_p)
                }
            };

            // Update rule: w_new = E{x * g(w^T * x)} - E{g'(w^T * x)} * w
            let e_g_prime: f64 = g_prime.iter().sum::<f64>() / n_samples as f64;

            for i in 0..n_components {
                let e_xg: f64 =
                    (0..n_samples).map(|t| whitened[i][t] * g[t]).sum::<f64>() / n_samples as f64;

                w[p][i] = e_xg - e_g_prime * w[p][i];
            }
        }

        // Orthogonalize
        orthogonalize(&mut w);

        // Check convergence
        let mut max_diff = 0.0f64;
        for p in 0..n_components {
            let dot: f64 = (0..n_components)
                .map(|i| w[p][i] * w_old[p][i])
                .sum::<f64>();
            let diff = 1.0 - dot.abs();
            max_diff = max_diff.max(diff);
        }

        if max_diff < tol {
            break;
        }
    }

    Ok(w)
}

/// Orthogonalize matrix using symmetric decorrelation
fn orthogonalize(w: &mut [Vec<f64>]) {
    let n = w.len();

    // W = W * (W^T * W)^(-1/2)
    // First compute W^T * W
    let mut wtw = vec![vec![0.0; n]; n];
    for i in 0..n {
        for j in 0..n {
            wtw[i][j] = (0..n).map(|k| w[i][k] * w[j][k]).sum();
        }
    }

    // Eigendecomposition
    if let Ok((eigvecs, eigvals)) = symmetric_eigen(&wtw) {
        // Compute (W^T * W)^(-1/2) = V * D^(-1/2) * V^T
        let mut inv_sqrt = vec![vec![0.0; n]; n];
        for i in 0..n {
            for j in 0..n {
                for k in 0..n {
                    let scale = if eigvals[k] > 1e-10 {
                        1.0 / eigvals[k].sqrt()
                    } else {
                        0.0
                    };
                    inv_sqrt[i][j] += eigvecs[k][i] * scale * eigvecs[k][j];
                }
            }
        }

        // W = W * inv_sqrt
        let w_old = w.to_vec();
        for i in 0..n {
            for j in 0..n {
                w[i][j] = (0..n).map(|k| w_old[i][k] * inv_sqrt[k][j]).sum();
            }
        }
    }
}

/// Symmetric eigendecomposition using power iteration
fn symmetric_eigen(matrix: &[Vec<f64>]) -> SignalResult<(Vec<Vec<f64>>, Vec<f64>)> {
    let n = matrix.len();
    if n == 0 {
        return Ok((vec![], vec![]));
    }

    let mut eigenvectors = Vec::with_capacity(n);
    let mut eigenvalues = Vec::with_capacity(n);
    let mut work = matrix.to_vec();

    for _ in 0..n {
        // Power iteration
        let mut v: Vec<f64> = (0..n).map(|i| ((i + 1) as f64).sin()).collect();
        let mut norm: f64 = v.iter().map(|&x| x * x).sum::<f64>().sqrt();
        for x in &mut v {
            *x /= norm;
        }

        let mut lambda = 0.0;
        for _ in 0..100 {
            // w = A * v
            let w: Vec<f64> = (0..n)
                .map(|i| (0..n).map(|j| work[i][j] * v[j]).sum())
                .collect();

            lambda = (0..n).map(|i| v[i] * w[i]).sum();
            norm = w.iter().map(|&x| x * x).sum::<f64>().sqrt();

            if norm < 1e-15 {
                break;
            }

            let diff: f64 = (0..n).map(|i| (w[i] / norm - v[i]).abs()).sum();

            v = w.iter().map(|&x| x / norm).collect();

            if diff < 1e-10 {
                break;
            }
        }

        if lambda.abs() < 1e-15 {
            break;
        }

        eigenvalues.push(lambda);
        eigenvectors.push(v.clone());

        // Deflate
        for i in 0..n {
            for j in 0..n {
                work[i][j] -= lambda * v[i] * v[j];
            }
        }
    }

    Ok((eigenvectors, eigenvalues))
}

/// Matrix multiplication
fn mat_mult(a: &[Vec<f64>], b: &[Vec<f64>]) -> Vec<Vec<f64>> {
    let m = a.len();
    let n = b[0].len();
    let k = b.len();

    let mut result = vec![vec![0.0; n]; m];
    for i in 0..m {
        for j in 0..n {
            for l in 0..k {
                result[i][j] += a[i][l] * b[l][j];
            }
        }
    }
    result
}

/// Pseudo-inverse using direct computation for square-ish matrices
fn pseudo_inverse(matrix: &[Vec<f64>]) -> SignalResult<Vec<Vec<f64>>> {
    let m = matrix.len();
    let n = matrix[0].len();

    // For the mixing matrix case (m x n where m >= n), compute (A^T A)^-1 A^T
    // Compute A^T * A
    let mut ata = vec![vec![0.0; n]; n];
    for i in 0..n {
        for j in 0..n {
            ata[i][j] = (0..m).map(|k| matrix[k][i] * matrix[k][j]).sum();
        }
    }

    // Regularize for numerical stability
    let trace: f64 = (0..n).map(|i| ata[i][i]).sum();
    let reg = 1e-6 * trace / n as f64;
    for i in 0..n {
        ata[i][i] += reg;
    }

    // Invert A^T * A using Cholesky-like method (symmetric positive definite)
    let ata_inv = invert_symmetric(&ata)?;

    // Compute (A^T A)^-1 * A^T
    let mut pinv = vec![vec![0.0; m]; n];
    for i in 0..n {
        for j in 0..m {
            for k in 0..n {
                pinv[i][j] += ata_inv[i][k] * matrix[j][k];
            }
        }
    }

    Ok(pinv)
}

/// Invert a symmetric positive definite matrix
fn invert_symmetric(matrix: &[Vec<f64>]) -> SignalResult<Vec<Vec<f64>>> {
    let n = matrix.len();

    // LDL decomposition
    let mut l = vec![vec![0.0; n]; n];
    let mut d = vec![0.0; n];

    for i in 0..n {
        // Compute D[i]
        let mut sum = matrix[i][i];
        for k in 0..i {
            sum -= l[i][k] * l[i][k] * d[k];
        }
        d[i] = sum;

        if d[i].abs() < 1e-15 {
            d[i] = 1e-10; // Regularize
        }

        l[i][i] = 1.0;

        // Compute L[j][i] for j > i
        for j in (i + 1)..n {
            let mut sum = matrix[j][i];
            for k in 0..i {
                sum -= l[j][k] * l[i][k] * d[k];
            }
            l[j][i] = sum / d[i];
        }
    }

    // Invert L
    let mut l_inv = vec![vec![0.0; n]; n];
    for i in 0..n {
        l_inv[i][i] = 1.0;
        for j in (i + 1)..n {
            let mut sum = 0.0;
            for k in i..j {
                sum -= l[j][k] * l_inv[k][i];
            }
            l_inv[j][i] = sum;
        }
    }

    // Compute A^-1 = L^-T * D^-1 * L^-1
    let mut result = vec![vec![0.0; n]; n];
    for i in 0..n {
        for j in 0..n {
            for k in 0..n {
                result[i][j] += l_inv[k][i] * (1.0 / d[k]) * l_inv[k][j];
            }
        }
    }

    Ok(result)
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::f64::consts::PI;

    fn create_test_data() -> Vec<Vec<f64>> {
        let n = 1000;
        let sfreq = 100.0;

        // Two independent source signals
        let s1: Vec<f64> = (0..n)
            .map(|i| (2.0 * PI * 3.0 * i as f64 / sfreq).sin())
            .collect();

        let s2: Vec<f64> = (0..n)
            .map(|i| {
                let t = i as f64 / sfreq;
                if (t * 2.0) as usize % 2 == 0 {
                    1.0
                } else {
                    -1.0
                }
            })
            .collect();

        // Mixing matrix
        let a = vec![vec![0.8, 0.6], vec![0.4, 0.9], vec![0.7, 0.3]];

        // Mixed signals: X = A * S
        let mut data = vec![vec![0.0; n]; 3];
        for ch in 0..3 {
            for t in 0..n {
                data[ch][t] = a[ch][0] * s1[t] + a[ch][1] * s2[t];
            }
        }

        data
    }

    #[test]
    fn test_ica_fit() {
        let data = create_test_data();

        let ica = Ica::fit(&data, Some(2), IcaMethod::FastIcaLogcosh, 200, 1e-4);
        assert!(ica.is_ok());

        let ica = ica.unwrap();
        assert_eq!(ica.n_components(), 2);
    }

    #[test]
    fn test_ica_transform() {
        let data = create_test_data();
        let ica = Ica::fit(&data, Some(2), IcaMethod::FastIcaLogcosh, 200, 1e-4).unwrap();

        let sources = ica.transform(&data).unwrap();
        assert_eq!(sources.len(), 2);
        assert_eq!(sources[0].len(), 1000);
    }

    #[test]
    fn test_ica_inverse() {
        let data = create_test_data();
        let ica = Ica::fit(&data, Some(2), IcaMethod::FastIcaLogcosh, 200, 1e-4).unwrap();

        let sources = ica.transform(&data).unwrap();
        let reconstructed = ica.inverse_transform(&sources).unwrap();

        // Check basic structure
        assert_eq!(reconstructed.len(), 3);
        assert_eq!(reconstructed[0].len(), 1000);

        // Verify that ICA decomposition and reconstruction work
        // (exact reconstruction is limited by numerical precision and component number)
        // The key test is that the inverse_transform doesn't crash and produces valid output
        for ch in 0..3 {
            // Check no NaN or Inf
            assert!(
                reconstructed[ch].iter().all(|x| x.is_finite()),
                "Reconstruction contains non-finite values in channel {}",
                ch
            );
        }
    }

    #[test]
    fn test_ica_apply() {
        let data = create_test_data();
        let ica = Ica::fit(&data, Some(2), IcaMethod::FastIcaLogcosh, 200, 1e-4).unwrap();

        // Remove first component
        let cleaned = ica.apply(&data, &[0]).unwrap();
        assert_eq!(cleaned.len(), 3);

        // Data should be different from original
        let diff: f64 = data[0]
            .iter()
            .zip(&cleaned[0])
            .map(|(&a, &b)| (a - b).abs())
            .sum();
        assert!(diff > 0.0);
    }

    #[test]
    fn test_component_kurtosis() {
        let data = create_test_data();
        let ica = Ica::fit(&data, Some(2), IcaMethod::FastIcaLogcosh, 200, 1e-4).unwrap();

        let kurtosis = ica.component_kurtosis(&data).unwrap();
        assert_eq!(kurtosis.len(), 2);
    }

    #[test]
    fn test_find_correlated() {
        let data = create_test_data();
        let ica = Ica::fit(&data, Some(2), IcaMethod::FastIcaLogcosh, 200, 1e-4).unwrap();

        // Use first channel as reference
        let reference = data[0].clone();
        let correlations = ica.find_correlated_components(&data, &reference).unwrap();

        assert_eq!(correlations.len(), 2);
        // All correlations should be bounded
        assert!(correlations.iter().all(|&c| c.abs() <= 1.0 + 1e-6));
    }

    #[test]
    fn test_ica_methods() {
        let data = create_test_data();

        // Test all methods
        for method in [
            IcaMethod::FastIcaLogcosh,
            IcaMethod::FastIcaExp,
            IcaMethod::FastIcaCube,
        ] {
            let ica = Ica::fit(&data, Some(2), method, 200, 1e-4);
            assert!(ica.is_ok(), "ICA {:?} failed", method);
        }
    }
}