//! Signal Space Projection (SSP) for artifact removal.
//!
//! SSP is a spatial filtering technique that removes artifacts by projecting
//! out directions in sensor space that contain artifact activity.
//!
//! ## Mathematical Background
//!
//! Given artifact data matrix A [n_channels x n_samples], we compute the
//! covariance C = A * A^T and find its principal components via SVD.
//!
//! The projection operator is: P = I - Σ(u_i * u_i^T)
//! where u_i are the artifact subspace vectors to remove.
//!
//! Clean data is obtained by: M_clean = P * M

use crate::{SignalError, SignalResult};

/// SSP projector for artifact removal
#[derive(Debug, Clone)]
pub struct SspProjector {
    /// Projection vectors [n_projectors x n_channels]
    /// Each row is a unit vector defining an artifact direction
    vectors: Vec<Vec<f64>>,
    /// Whether each projector is active
    active: Vec<bool>,
    /// Explained variance by each projector
    explained_var: Vec<f64>,
    /// Names/descriptions of projectors
    names: Vec<String>,
    /// Number of channels
    n_channels: usize,
}

impl SspProjector {
    /// Create an SSP projector from artifact epochs
    ///
    /// # Arguments
    /// * `epochs` - Artifact epochs [n_epochs][n_channels][n_samples]
    /// * `n_components` - Number of projectors to compute
    /// * `name` - Name for this projector set (e.g., "EOG", "ECG")
    ///
    /// # Returns
    /// SSP projector with the requested number of components
    pub fn from_epochs(
        epochs: &[Vec<Vec<f64>>],
        n_components: usize,
        name: &str,
    ) -> SignalResult<Self> {
        if epochs.is_empty() {
            return Err(SignalError::InvalidLength("No epochs provided".to_string()));
        }

        let n_channels = epochs[0].len();
        if n_channels == 0 {
            return Err(SignalError::InvalidLength("Empty channels".to_string()));
        }

        // Concatenate all epochs into a data matrix [n_channels x total_samples]
        let total_samples: usize = epochs.iter().map(|e| e[0].len()).sum();
        let mut data = vec![vec![0.0; total_samples]; n_channels];

        let mut offset = 0;
        for epoch in epochs {
            let n_times = epoch[0].len();
            for ch in 0..n_channels {
                for t in 0..n_times {
                    data[ch][offset + t] = epoch[ch][t];
                }
            }
            offset += n_times;
        }

        // Compute covariance matrix: C = (1/n) * X * X^T
        let mut cov = vec![vec![0.0; n_channels]; n_channels];
        for i in 0..n_channels {
            for j in i..n_channels {
                let mut sum = 0.0;
                for t in 0..total_samples {
                    sum += data[i][t] * data[j][t];
                }
                cov[i][j] = sum / total_samples as f64;
                cov[j][i] = cov[i][j];
            }
        }

        // Compute SVD of covariance matrix
        let (eigenvectors, eigenvalues) = symmetric_eigen(&cov)?;

        // Total variance
        let total_var: f64 = eigenvalues.iter().sum();

        // Extract top n_components
        let n_proj = n_components.min(n_channels);
        let mut vectors = Vec::with_capacity(n_proj);
        let mut explained_var = Vec::with_capacity(n_proj);
        let mut names = Vec::with_capacity(n_proj);
        let mut active = Vec::with_capacity(n_proj);

        for i in 0..n_proj {
            vectors.push(eigenvectors[i].clone());
            explained_var.push(eigenvalues[i] / total_var);
            names.push(format!("{}-{}", name, i + 1));
            active.push(true);
        }

        Ok(Self {
            vectors,
            active,
            explained_var,
            names,
            n_channels,
        })
    }

    /// Create SSP projector from EOG data
    ///
    /// # Arguments
    /// * `data` - Continuous data [n_channels][n_samples]
    /// * `eog_channels` - Indices of EOG channels
    /// * `n_components` - Number of projectors (typically 1-2)
    pub fn from_eog(
        data: &[Vec<f64>],
        eog_channels: &[usize],
        n_components: usize,
    ) -> SignalResult<Self> {
        if data.is_empty() {
            return Err(SignalError::InvalidLength("Empty data".to_string()));
        }

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

        // Create artifact epochs from EOG threshold crossings
        let mut artifact_epochs = Vec::new();

        for &eog_ch in eog_channels {
            if eog_ch >= n_channels {
                return Err(SignalError::InvalidParameters(format!(
                    "EOG channel {eog_ch} out of range"
                )));
            }

            // Find threshold (use robust estimate: 5 * median absolute deviation)
            let eog_data = &data[eog_ch];
            let threshold = compute_robust_threshold(eog_data);

            // Find peaks above threshold
            let peaks = find_peaks(eog_data, threshold);

            // Extract epochs around each peak (±100 samples or available)
            let half_window = 100.min(n_samples / 2);
            for peak in peaks {
                let start = peak.saturating_sub(half_window);
                let end = (peak + half_window).min(n_samples);

                let epoch: Vec<Vec<f64>> = data.iter().map(|ch| ch[start..end].to_vec()).collect();
                artifact_epochs.push(epoch);
            }
        }

        if artifact_epochs.is_empty() {
            // Create a projector from the raw EOG channel correlation
            return Self::from_eog_correlation(data, eog_channels, n_components);
        }

        Self::from_epochs(&artifact_epochs, n_components, "EOG")
    }

    /// Create SSP from EOG channel correlation (fallback method)
    fn from_eog_correlation(
        data: &[Vec<f64>],
        eog_channels: &[usize],
        n_components: usize,
    ) -> SignalResult<Self> {
        let n_channels = data.len();
        let n_samples = data[0].len();

        // Compute correlation of each channel with EOG channels
        let mut vectors = Vec::new();
        let mut explained_var = Vec::new();

        for (comp_idx, &eog_ch) in eog_channels.iter().take(n_components).enumerate() {
            let eog_data = &data[eog_ch];
            let eog_var: f64 = eog_data.iter().map(|&x| x * x).sum::<f64>() / n_samples as f64;

            if eog_var < 1e-15 {
                continue;
            }

            // Compute correlation vector
            let mut corr_vec = vec![0.0; n_channels];
            for ch in 0..n_channels {
                let cov: f64 = (0..n_samples)
                    .map(|t| data[ch][t] * eog_data[t])
                    .sum::<f64>()
                    / n_samples as f64;
                corr_vec[ch] = cov / eog_var.sqrt();
            }

            // Normalize
            let norm: f64 = corr_vec.iter().map(|&x| x * x).sum::<f64>().sqrt();
            if norm > 1e-10 {
                for v in &mut corr_vec {
                    *v /= norm;
                }
                vectors.push(corr_vec);
                explained_var.push(1.0 / (comp_idx + 1) as f64);
            }
        }

        let n_proj = vectors.len();
        Ok(Self {
            vectors,
            active: vec![true; n_proj],
            explained_var,
            names: (0..n_proj).map(|i| format!("EOG-{}", i + 1)).collect(),
            n_channels,
        })
    }

    /// Create SSP projector from ECG data
    ///
    /// # Arguments
    /// * `data` - Continuous data [n_channels][n_samples]
    /// * `ecg_channel` - Index of ECG channel
    /// * `n_components` - Number of projectors (typically 1-3)
    pub fn from_ecg(
        data: &[Vec<f64>],
        ecg_channel: usize,
        n_components: usize,
    ) -> SignalResult<Self> {
        if data.is_empty() {
            return Err(SignalError::InvalidLength("Empty data".to_string()));
        }

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

        if ecg_channel >= n_channels {
            return Err(SignalError::InvalidParameters(format!(
                "ECG channel {ecg_channel} out of range"
            )));
        }

        let ecg_data = &data[ecg_channel];

        // Find R-peaks using simple threshold detection
        let threshold = compute_robust_threshold(ecg_data);
        let r_peaks = find_peaks(ecg_data, threshold);

        if r_peaks.len() < 3 {
            return Err(SignalError::InvalidLength(
                "Not enough ECG peaks detected".to_string(),
            ));
        }

        // Extract epochs around each R-peak (±50 samples)
        let half_window = 50.min(n_samples / 2);
        let mut artifact_epochs = Vec::new();

        for peak in r_peaks {
            let start = peak.saturating_sub(half_window);
            let end = (peak + half_window).min(n_samples);

            let epoch: Vec<Vec<f64>> = data.iter().map(|ch| ch[start..end].to_vec()).collect();
            artifact_epochs.push(epoch);
        }

        Self::from_epochs(&artifact_epochs, n_components, "ECG")
    }

    /// Apply the SSP projector to data
    ///
    /// # Arguments
    /// * `data` - Data to clean [n_channels][n_samples]
    ///
    /// # Returns
    /// Cleaned data with artifacts removed
    pub fn apply(&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();

        // Compute projection matrix: P = I - Σ(u_i * u_i^T)
        // For efficiency, we apply it directly: y = x - Σ((u_i^T * x) * u_i)
        let mut result = data.to_vec();

        for (i, vec) in self.vectors.iter().enumerate() {
            if !self.active[i] {
                continue;
            }

            for t in 0..n_samples {
                // Compute dot product: u^T * x
                let dot: f64 = (0..self.n_channels).map(|ch| vec[ch] * data[ch][t]).sum();

                // Subtract: x = x - (u^T * x) * u
                for ch in 0..self.n_channels {
                    result[ch][t] -= dot * vec[ch];
                }
            }
        }

        Ok(result)
    }

    /// Apply SSP to epoched data
    ///
    /// # Arguments
    /// * `epochs` - Epoched data [n_epochs][n_channels][n_samples]
    ///
    /// # Returns
    /// Cleaned epochs
    pub fn apply_epochs(&self, epochs: &[Vec<Vec<f64>>]) -> SignalResult<Vec<Vec<Vec<f64>>>> {
        epochs.iter().map(|epoch| self.apply(epoch)).collect()
    }

    /// Activate or deactivate a projector
    pub fn set_active(&mut self, index: usize, active: bool) -> SignalResult<()> {
        if index >= self.vectors.len() {
            return Err(SignalError::InvalidParameters(format!(
                "Projector index {index} out of range"
            )));
        }
        self.active[index] = active;
        Ok(())
    }

    /// Get number of projectors
    pub fn n_projectors(&self) -> usize {
        self.vectors.len()
    }

    /// Get number of active projectors
    pub fn n_active(&self) -> usize {
        self.active.iter().filter(|&&a| a).count()
    }

    /// Get projector vectors
    pub fn vectors(&self) -> &[Vec<f64>] {
        &self.vectors
    }

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

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

    /// Check if projector is active
    pub fn is_active(&self, index: usize) -> bool {
        self.active.get(index).copied().unwrap_or(false)
    }

    /// Compute the projection matrix P = I - Σ(u_i * u_i^T)
    pub fn compute_projection_matrix(&self) -> Vec<Vec<f64>> {
        let n = self.n_channels;
        let mut proj = vec![vec![0.0; n]; n];

        // Start with identity
        for i in 0..n {
            proj[i][i] = 1.0;
        }

        // Subtract outer products of active projectors
        for (i, vec) in self.vectors.iter().enumerate() {
            if !self.active[i] {
                continue;
            }
            for row in 0..n {
                for col in 0..n {
                    proj[row][col] -= vec[row] * vec[col];
                }
            }
        }

        proj
    }

    /// Merge multiple SSP projectors
    pub fn merge(projectors: &[&SspProjector]) -> SignalResult<Self> {
        if projectors.is_empty() {
            return Err(SignalError::InvalidParameters(
                "No projectors to merge".to_string(),
            ));
        }

        let n_channels = projectors[0].n_channels;
        for p in projectors {
            if p.n_channels != n_channels {
                return Err(SignalError::InvalidParameters(
                    "Channel count mismatch".to_string(),
                ));
            }
        }

        let mut vectors = Vec::new();
        let mut active = Vec::new();
        let mut explained_var = Vec::new();
        let mut names = Vec::new();

        for p in projectors {
            vectors.extend(p.vectors.clone());
            active.extend(p.active.clone());
            explained_var.extend(p.explained_var.clone());
            names.extend(p.names.clone());
        }

        Ok(Self {
            vectors,
            active,
            explained_var,
            names,
            n_channels,
        })
    }
}

/// Compute robust threshold using median absolute deviation
fn compute_robust_threshold(data: &[f64]) -> f64 {
    let mut sorted: Vec<f64> = data.iter().map(|&x| x.abs()).collect();
    sorted.sort_by(|a, b| a.partial_cmp(b).unwrap());

    let median = if sorted.len().is_multiple_of(2) {
        f64::midpoint(sorted[sorted.len() / 2 - 1], sorted[sorted.len() / 2])
    } else {
        sorted[sorted.len() / 2]
    };

    // MAD = median(|x - median|)
    let mut deviations: Vec<f64> = data.iter().map(|&x| (x.abs() - median).abs()).collect();
    deviations.sort_by(|a, b| a.partial_cmp(b).unwrap());

    let mad = if deviations.len().is_multiple_of(2) {
        f64::midpoint(
            deviations[deviations.len() / 2 - 1],
            deviations[deviations.len() / 2],
        )
    } else {
        deviations[deviations.len() / 2]
    };

    // 5 * MAD is a robust threshold
    5.0 * mad.max(1e-10)
}

/// Find peaks above threshold with minimum separation
fn find_peaks(data: &[f64], threshold: f64) -> Vec<usize> {
    let min_separation = 50; // Minimum samples between peaks
    let mut peaks = Vec::new();
    let mut last_peak: Option<usize> = None;

    for i in 1..data.len().saturating_sub(1) {
        if data[i] > threshold && data[i] > data[i - 1] && data[i] > data[i + 1] {
            if let Some(last) = last_peak {
                if i - last >= min_separation {
                    peaks.push(i);
                    last_peak = Some(i);
                }
            } else {
                peaks.push(i);
                last_peak = Some(i);
            }
        }
    }

    peaks
}

/// Compute eigendecomposition of symmetric matrix 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 = matrix.to_vec();

    for _ in 0..n {
        // Power iteration to find largest eigenvalue/vector
        let (eigval, eigvec) = power_iteration(&work_matrix, 100)?;

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

        eigenvalues.push(eigval);
        eigenvectors.push(eigvec.clone());

        // Deflate: A = A - λ * v * v^T
        for i in 0..n {
            for j in 0..n {
                work_matrix[i][j] -= eigval * eigvec[i] * eigvec[j];
            }
        }
    }

    Ok((eigenvectors, eigenvalues))
}

/// Power iteration to find dominant eigenvalue/vector
fn power_iteration(matrix: &[Vec<f64>], max_iter: usize) -> SignalResult<(f64, Vec<f64>)> {
    let n = matrix.len();
    if n == 0 {
        return Err(SignalError::InvalidLength("Empty matrix".to_string()));
    }

    // Initialize with random-ish vector
    let mut v: Vec<f64> = (0..n).map(|i| ((i + 1) as f64).sin()).collect();

    // Normalize
    let mut norm: f64 = v.iter().map(|&x| x * x).sum::<f64>().sqrt();
    for x in &mut v {
        *x /= norm;
    }

    let mut eigenvalue = 0.0;

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

        // Rayleigh quotient: λ = v^T * A * v
        eigenvalue = (0..n).map(|i| v[i] * w[i]).sum();

        // Normalize w
        norm = w.iter().map(|&x| x * x).sum::<f64>().sqrt();
        if norm < 1e-15 {
            break;
        }

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

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

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

    Ok((eigenvalue, v))
}

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

    #[test]
    fn test_ssp_from_epochs() {
        // Create simple artifact epochs
        let epochs = vec![
            vec![
                vec![1.0, 2.0, 3.0, 2.0, 1.0],
                vec![0.5, 1.0, 1.5, 1.0, 0.5],
                vec![0.2, 0.4, 0.6, 0.4, 0.2],
            ],
            vec![
                vec![1.1, 2.1, 3.1, 2.1, 1.1],
                vec![0.55, 1.05, 1.55, 1.05, 0.55],
                vec![0.22, 0.42, 0.62, 0.42, 0.22],
            ],
        ];

        let ssp = SspProjector::from_epochs(&epochs, 2, "test").unwrap();

        assert!(ssp.n_projectors() >= 1);
        assert_eq!(ssp.n_active(), ssp.n_projectors());
    }

    #[test]
    fn test_ssp_apply() {
        // Create a known projector
        let ssp = SspProjector {
            vectors: vec![vec![1.0, 0.0, 0.0]], // Project out first channel
            active: vec![true],
            explained_var: vec![0.5],
            names: vec!["test-1".to_string()],
            n_channels: 3,
        };

        let data = vec![
            vec![1.0, 2.0, 3.0],
            vec![4.0, 5.0, 6.0],
            vec![7.0, 8.0, 9.0],
        ];

        let cleaned = ssp.apply(&data).unwrap();

        // First channel should be zeroed
        assert!(cleaned[0].iter().all(|&x| x.abs() < 1e-10));
        // Other channels unchanged
        assert_eq!(cleaned[1], vec![4.0, 5.0, 6.0]);
        assert_eq!(cleaned[2], vec![7.0, 8.0, 9.0]);
    }

    #[test]
    fn test_ssp_projection_matrix() {
        let ssp = SspProjector {
            vectors: vec![vec![1.0, 0.0], vec![0.0, 1.0]],
            active: vec![true, false],
            explained_var: vec![0.5, 0.3],
            names: vec!["p1".to_string(), "p2".to_string()],
            n_channels: 2,
        };

        let proj = ssp.compute_projection_matrix();

        // Only first projector is active, so P = I - u1*u1^T
        // u1 = [1, 0], so P = [[0, 0], [0, 1]]
        assert!((proj[0][0] - 0.0).abs() < 1e-10);
        assert!((proj[0][1] - 0.0).abs() < 1e-10);
        assert!((proj[1][0] - 0.0).abs() < 1e-10);
        assert!((proj[1][1] - 1.0).abs() < 1e-10);
    }

    #[test]
    fn test_set_active() {
        let mut ssp = SspProjector {
            vectors: vec![vec![1.0, 0.0]],
            active: vec![true],
            explained_var: vec![0.5],
            names: vec!["p1".to_string()],
            n_channels: 2,
        };

        assert!(ssp.is_active(0));
        ssp.set_active(0, false).unwrap();
        assert!(!ssp.is_active(0));
        assert_eq!(ssp.n_active(), 0);
    }

    #[test]
    fn test_merge_projectors() {
        let p1 = SspProjector {
            vectors: vec![vec![1.0, 0.0, 0.0]],
            active: vec![true],
            explained_var: vec![0.5],
            names: vec!["EOG-1".to_string()],
            n_channels: 3,
        };

        let p2 = SspProjector {
            vectors: vec![vec![0.0, 1.0, 0.0]],
            active: vec![true],
            explained_var: vec![0.3],
            names: vec!["ECG-1".to_string()],
            n_channels: 3,
        };

        let merged = SspProjector::merge(&[&p1, &p2]).unwrap();

        assert_eq!(merged.n_projectors(), 2);
        assert_eq!(merged.names()[0], "EOG-1");
        assert_eq!(merged.names()[1], "ECG-1");
    }

    #[test]
    fn test_power_iteration() {
        // Diagonal matrix - eigenvalues are diagonal elements
        let matrix = vec![vec![3.0, 0.0], vec![0.0, 1.0]];

        let (eigval, eigvec) = power_iteration(&matrix, 100).unwrap();

        // Largest eigenvalue should be 3
        assert!((eigval - 3.0).abs() < 1e-6);
        // Corresponding eigenvector should be [1, 0]
        assert!((eigvec[0].abs() - 1.0).abs() < 1e-6);
        assert!(eigvec[1].abs() < 1e-6);
    }

    #[test]
    fn test_find_peaks() {
        let data = vec![0.0, 1.0, 5.0, 2.0, 0.0, 0.0, 0.0];
        let peaks = find_peaks(&data, 3.0);

        assert_eq!(peaks.len(), 1);
        assert_eq!(peaks[0], 2);
    }
}