//! GPU-accelerated LCMV beamformer for real-time source localization.
//!
//! Provides sub-millisecond beamformer computation using pre-computed spatial filters.

use crate::error::{RealtimeError, RealtimeResult};
use nalgebra::{DMatrix, DVector};
use std::sync::Arc;

/// GPU beamformer configuration
#[derive(Debug, Clone)]
pub struct GpuBeamformerConfig {
    /// Number of sensor channels
    pub n_channels: usize,
    /// Number of source locations
    pub n_sources: usize,
    /// Whether sources have free orientation (3 components per source)
    pub free_orientation: bool,
    /// Regularization parameter (fraction of trace)
    pub regularization: f64,
    /// Power normalization enabled
    pub normalize_power: bool,
}

impl Default for GpuBeamformerConfig {
    fn default() -> Self {
        Self {
            n_channels: 64,
            n_sources: 1000,
            free_orientation: false,
            regularization: 0.05,
            normalize_power: true,
        }
    }
}

/// Pre-computed beamformer weights for GPU processing
#[derive(Debug, Clone)]
pub struct BeamformerWeights {
    /// Spatial filters [n_source_components x n_channels]
    pub weights: DMatrix<f64>,
    /// Power normalization factors [n_source_components]
    pub power_norm: DVector<f64>,
    /// Source indices
    pub source_indices: Vec<usize>,
    /// Number of source components (n_sources * 3 if free orientation)
    pub n_components: usize,
}

impl BeamformerWeights {
    /// Create from pre-computed weights
    pub fn new(
        weights: DMatrix<f64>,
        power_norm: Option<DVector<f64>>,
        source_indices: Vec<usize>,
    ) -> Self {
        let n_components = weights.nrows();
        let power_norm = power_norm.unwrap_or_else(|| DVector::from_element(n_components, 1.0));

        Self {
            weights,
            power_norm,
            source_indices,
            n_components,
        }
    }

    /// Number of channels
    pub fn n_channels(&self) -> usize {
        self.weights.ncols()
    }

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

/// GPU-accelerated LCMV beamformer
#[derive(Debug)]
pub struct GpuBeamformer {
    /// Configuration
    config: GpuBeamformerConfig,
    /// Pre-computed beamformer weights
    weights: Arc<BeamformerWeights>,
    /// Cached covariance inverse (for weight updates)
    cov_inv: Option<DMatrix<f64>>,
}

impl GpuBeamformer {
    /// Create a new GPU beamformer from pre-computed weights
    ///
    /// # Arguments
    /// * `config` - Beamformer configuration
    /// * `weights` - Pre-computed spatial filter weights
    pub fn new(config: GpuBeamformerConfig, weights: BeamformerWeights) -> RealtimeResult<Self> {
        // Validate dimensions
        if weights.n_channels() != config.n_channels {
            return Err(RealtimeError::DimensionMismatch(format!(
                "Weights have {} channels, config expects {}",
                weights.n_channels(),
                config.n_channels
            )));
        }

        Ok(Self {
            config,
            weights: Arc::new(weights),
            cov_inv: None,
        })
    }

    /// Create a beamformer from gain matrix and covariance
    ///
    /// # Arguments
    /// * `config` - Beamformer configuration
    /// * `gain` - Forward model gain matrix [n_channels x n_sources]
    /// * `data_cov` - Data covariance matrix [n_channels x n_channels]
    pub fn from_forward_model(
        config: GpuBeamformerConfig,
        gain: &DMatrix<f64>,
        data_cov: &DMatrix<f64>,
    ) -> RealtimeResult<Self> {
        let n_channels = gain.nrows();
        let n_source_cols = gain.ncols();

        if data_cov.nrows() != n_channels || data_cov.ncols() != n_channels {
            return Err(RealtimeError::DimensionMismatch(
                "Covariance matrix dimension mismatch".to_string(),
            ));
        }

        // Regularize and invert covariance
        let mut c = data_cov.clone();
        let trace: f64 = (0..n_channels).map(|i| c[(i, i)]).sum();
        let reg_val = config.regularization * trace / n_channels as f64;

        for i in 0..n_channels {
            c[(i, i)] += reg_val;
        }

        let c_inv = Self::invert_matrix(&c)?;

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

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

            if norm > 1e-30 {
                for ch in 0..n_channels {
                    weights[(src, ch)] = c_inv_g[ch] / norm;
                }
                power_norm[src] = 1.0 / norm.sqrt();
            }
        }

        let source_indices: Vec<usize> = if config.free_orientation {
            (0..n_source_cols / 3).collect()
        } else {
            (0..n_source_cols).collect()
        };

        let bf_weights = BeamformerWeights::new(weights, Some(power_norm), source_indices);

        let mut bf = Self::new(config, bf_weights)?;
        bf.cov_inv = Some(c_inv);

        Ok(bf)
    }

    /// Apply 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>]) -> RealtimeResult<Vec<Vec<f64>>> {
        if data.len() != self.config.n_channels {
            return Err(RealtimeError::DimensionMismatch(format!(
                "Expected {} channels, got {}",
                self.config.n_channels,
                data.len()
            )));
        }

        if data.is_empty() || data[0].is_empty() {
            return Ok(vec![vec![]; self.weights.n_components]);
        }

        let n_times = data[0].len();
        let n_comp = self.weights.n_components;

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

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

        // Apply power normalization if enabled
        let output: Vec<Vec<f64>> = if self.config.normalize_power {
            (0..n_comp)
                .map(|i| {
                    let norm = self.weights.power_norm[i];
                    (0..n_times).map(|j| source_mat[(i, j)] * norm).collect()
                })
                .collect()
        } else {
            (0..n_comp)
                .map(|i| (0..n_times).map(|j| source_mat[(i, j)]).collect())
                .collect()
        };

        // If free orientation, combine xyz components into power
        if self.config.free_orientation {
            let n_sources = n_comp / 3;
            let combined: Vec<Vec<f64>> = (0..n_sources)
                .map(|s| {
                    let x = &output[s * 3];
                    let y = &output[s * 3 + 1];
                    let z = &output[s * 3 + 2];
                    (0..n_times)
                        .map(|t| (x[t].powi(2) + y[t].powi(2) + z[t].powi(2)).sqrt())
                        .collect()
                })
                .collect();
            Ok(combined)
        } else {
            Ok(output)
        }
    }

    /// Apply beamformer using nalgebra matrix (faster for batch processing)
    pub fn apply_matrix(&self, data: &DMatrix<f64>) -> RealtimeResult<DMatrix<f64>> {
        if data.nrows() != self.config.n_channels {
            return Err(RealtimeError::DimensionMismatch(format!(
                "Expected {} channels, got {}",
                self.config.n_channels,
                data.nrows()
            )));
        }

        // Apply spatial filters: S = W * M
        let mut source_mat = &self.weights.weights * data;

        // Apply power normalization
        if self.config.normalize_power {
            for i in 0..source_mat.nrows() {
                let norm = self.weights.power_norm[i];
                for j in 0..source_mat.ncols() {
                    source_mat[(i, j)] *= norm;
                }
            }
        }

        Ok(source_mat)
    }

    /// Compute Neural Activity Index (NAI) for each source
    ///
    /// NAI = (w^T * C_data * w) / (w^T * C_noise * w)
    pub fn compute_nai(
        &self,
        data_cov: &DMatrix<f64>,
        noise_cov: &DMatrix<f64>,
    ) -> RealtimeResult<Vec<f64>> {
        let n_comp = self.weights.n_components;
        let n_channels = self.config.n_channels;

        if data_cov.nrows() != n_channels || noise_cov.nrows() != n_channels {
            return Err(RealtimeError::DimensionMismatch(
                "Covariance dimension mismatch".to_string(),
            ));
        }

        let mut nai = Vec::with_capacity(n_comp);

        for src in 0..n_comp {
            let w = self.weights.weights.row(src);

            // w^T * C * w
            let compute_power = |cov: &DMatrix<f64>| -> f64 {
                (0..n_channels)
                    .map(|i| {
                        (0..n_channels)
                            .map(|j| w[i] * cov[(i, j)] * w[j])
                            .sum::<f64>()
                    })
                    .sum()
            };

            let signal_power = compute_power(data_cov);
            let noise_power = compute_power(noise_cov);

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

        Ok(nai)
    }

    /// Get the spatial filter weights
    pub fn weights(&self) -> &BeamformerWeights {
        &self.weights
    }

    /// Get configuration
    pub fn config(&self) -> &GpuBeamformerConfig {
        &self.config
    }

    /// Update beamformer weights with new covariance (adaptive beamforming)
    pub fn update_covariance(&mut self, new_cov: &DMatrix<f64>) -> RealtimeResult<()> {
        // This would require re-computing the weights
        // For real-time use, this should be done asynchronously
        Err(RealtimeError::NotInitialized(
            "Adaptive beamforming not yet implemented".to_string(),
        ))
    }

    /// Invert matrix using SVD
    fn invert_matrix(m: &DMatrix<f64>) -> RealtimeResult<DMatrix<f64>> {
        let n = m.nrows();
        let svd = m.clone().svd(true, true);

        let u = svd
            .u
            .ok_or_else(|| RealtimeError::Beamformer("SVD failed to compute U".to_string()))?;
        let vt = svd
            .v_t
            .ok_or_else(|| RealtimeError::Beamformer("SVD failed to compute V^T".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())
    }
}

/// Create a beamformer for a subset of sources (ROI-based)
pub fn create_roi_beamformer(
    full_beamformer: &GpuBeamformer,
    source_indices: &[usize],
) -> RealtimeResult<GpuBeamformer> {
    let full_weights = &full_beamformer.weights;
    let mut config = full_beamformer.config.clone();

    let n_channels = full_weights.n_channels();
    let n_selected = source_indices.len();

    // Extract weights for selected sources
    let mut selected_weights = DMatrix::zeros(n_selected, n_channels);
    let mut selected_norm = DVector::zeros(n_selected);

    for (i, &src_idx) in source_indices.iter().enumerate() {
        if src_idx >= full_weights.n_components {
            return Err(RealtimeError::DimensionMismatch(format!(
                "Source index {} out of range (max {})",
                src_idx,
                full_weights.n_components - 1
            )));
        }

        for ch in 0..n_channels {
            selected_weights[(i, ch)] = full_weights.weights[(src_idx, ch)];
        }
        selected_norm[i] = full_weights.power_norm[src_idx];
    }

    config.n_sources = n_selected;
    config.free_orientation = false;

    let roi_weights = BeamformerWeights::new(
        selected_weights,
        Some(selected_norm),
        source_indices.to_vec(),
    );

    GpuBeamformer::new(config, roi_weights)
}

/// Compute source power (envelope) from beamformer output
pub fn compute_source_power(sources: &[Vec<f64>], window_samples: usize) -> Vec<Vec<f64>> {
    sources
        .iter()
        .map(|source| {
            source
                .windows(window_samples)
                .map(|window| {
                    let sum_sq: f64 = window.iter().map(|x| x.powi(2)).sum();
                    (sum_sq / window_samples as f64).sqrt()
                })
                .collect()
        })
        .collect()
}

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

    fn create_test_gain() -> DMatrix<f64> {
        // Simple 4-channel, 3-source gain matrix
        DMatrix::from_row_slice(
            4,
            3,
            &[1.0, 0.5, 0.2, 0.5, 1.0, 0.5, 0.2, 0.5, 1.0, 0.1, 0.3, 0.6],
        )
    }

    fn create_test_covariance() -> DMatrix<f64> {
        // Simple identity-like covariance
        DMatrix::from_diagonal(&DVector::from_vec(vec![1.0, 1.0, 1.0, 1.0]))
    }

    #[test]
    fn test_beamformer_creation() {
        let config = GpuBeamformerConfig {
            n_channels: 4,
            n_sources: 3,
            free_orientation: false,
            regularization: 0.05,
            normalize_power: true,
        };

        let gain = create_test_gain();
        let cov = create_test_covariance();

        let bf = GpuBeamformer::from_forward_model(config, &gain, &cov).unwrap();

        assert_eq!(bf.weights.n_channels(), 4);
        assert_eq!(bf.weights.n_components(), 3);
    }

    #[test]
    fn test_beamformer_apply() {
        let config = GpuBeamformerConfig {
            n_channels: 4,
            n_sources: 3,
            free_orientation: false,
            regularization: 0.05,
            normalize_power: false,
        };

        let gain = create_test_gain();
        let cov = create_test_covariance();
        let bf = GpuBeamformer::from_forward_model(config, &gain, &cov).unwrap();

        // Test data: 4 channels, 10 time points
        let data: Vec<Vec<f64>> = (0..4)
            .map(|ch| (0..10).map(|t| (ch * 10 + t) as f64).collect())
            .collect();

        let sources = bf.apply(&data).unwrap();

        assert_eq!(sources.len(), 3);
        assert_eq!(sources[0].len(), 10);
    }

    #[test]
    fn test_apply_matrix() {
        let config = GpuBeamformerConfig {
            n_channels: 4,
            n_sources: 3,
            free_orientation: false,
            regularization: 0.05,
            normalize_power: false,
        };

        let gain = create_test_gain();
        let cov = create_test_covariance();
        let bf = GpuBeamformer::from_forward_model(config, &gain, &cov).unwrap();

        let data = DMatrix::from_fn(4, 10, |i, j| (i * 10 + j) as f64);
        let sources = bf.apply_matrix(&data).unwrap();

        assert_eq!(sources.nrows(), 3);
        assert_eq!(sources.ncols(), 10);
    }

    #[test]
    fn test_compute_nai() {
        let config = GpuBeamformerConfig {
            n_channels: 4,
            n_sources: 3,
            free_orientation: false,
            regularization: 0.05,
            normalize_power: false,
        };

        let gain = create_test_gain();
        let cov = create_test_covariance();
        let bf = GpuBeamformer::from_forward_model(config, &gain, &cov).unwrap();

        // Data cov with higher signal
        let data_cov = DMatrix::from_diagonal(&DVector::from_vec(vec![2.0, 2.0, 2.0, 2.0]));
        let noise_cov = DMatrix::from_diagonal(&DVector::from_vec(vec![1.0, 1.0, 1.0, 1.0]));

        let nai = bf.compute_nai(&data_cov, &noise_cov).unwrap();

        assert_eq!(nai.len(), 3);
        // NAI should be positive with higher signal than noise
        for n in &nai {
            assert!(*n >= 0.0);
        }
    }

    #[test]
    fn test_roi_beamformer() {
        let config = GpuBeamformerConfig {
            n_channels: 4,
            n_sources: 3,
            free_orientation: false,
            regularization: 0.05,
            normalize_power: false,
        };

        let gain = create_test_gain();
        let cov = create_test_covariance();
        let bf = GpuBeamformer::from_forward_model(config, &gain, &cov).unwrap();

        // Create ROI beamformer for sources 0 and 2
        let roi_bf = create_roi_beamformer(&bf, &[0, 2]).unwrap();

        assert_eq!(roi_bf.weights.n_components(), 2);
    }

    #[test]
    fn test_source_power() {
        let sources = vec![vec![1.0, 2.0, 3.0, 4.0, 5.0], vec![0.5, 1.0, 1.5, 2.0, 2.5]];

        let power = compute_source_power(&sources, 3);

        assert_eq!(power.len(), 2);
        assert_eq!(power[0].len(), 3); // 5 - 3 + 1 = 3 windows
    }
}