//! GPU-accelerated Non-negative Matrix Factorization
//!
//! Provides NMF decomposition using GPU acceleration when available,
//! with automatic fallback to CPU computation.

use crate::config::NmfProcessingConfig;
use crate::gpu::{BackendType, GpuBackend, select_backend};
use crate::nmf::NmfOutput;
use anyhow::{Result, bail};
use ndarray::Array2;
use tracing::{debug, info};

/// GPU-accelerated NMF processor
///
/// Uses CUDA on NVIDIA GPUs, Metal on Apple Silicon, or falls back to CPU.
pub struct GpuNmfProcessor {
    config: NmfProcessingConfig,
    backend: Box<dyn GpuBackend>,
}

impl GpuNmfProcessor {
    /// Create a new GPU NMF processor with the given configuration
    ///
    /// Automatically selects the best available GPU backend.
    pub fn new(config: NmfProcessingConfig) -> Result<Self> {
        let backend = select_backend()?;
        let device = backend.device_info();

        info!(
            "GpuNmfProcessor initialized with {} backend: {}",
            device.backend, device.name
        );

        Ok(Self { config, backend })
    }

    /// Create a processor with a specific backend
    pub fn with_backend(config: NmfProcessingConfig, backend: Box<dyn GpuBackend>) -> Self {
        Self { config, backend }
    }

    /// Get the backend type being used
    pub fn backend_type(&self) -> BackendType {
        self.backend.device_info().backend
    }

    /// Check if GPU acceleration is active (not CPU fallback)
    pub fn is_gpu_accelerated(&self) -> bool {
        matches!(
            self.backend.device_info().backend,
            BackendType::Cuda | BackendType::Metal
        )
    }

    /// Perform NMF decomposition: V ≈ W * H
    ///
    /// # Arguments
    /// * `v` - Input matrix (n_features x n_samples), e.g., (3 x n_pixels) for RGB/OD
    /// * `w_init` - Optional initial W matrix (stain vectors)
    ///
    /// # Returns
    /// NMF output containing W and H matrices
    pub fn decompose(&self, v: &Array2<f32>, w_init: Option<&Array2<f32>>) -> Result<NmfOutput> {
        let (n_features, n_samples) = v.dim();
        let k = self.config.n_components;

        info!(
            "GPU NMF: {}x{} matrix, {} components, max {} iterations ({})",
            n_features,
            n_samples,
            k,
            self.config.max_iterations,
            self.backend.name()
        );

        if n_features == 0 || n_samples == 0 {
            bail!("Input matrix must not be empty");
        }
        if k > n_features || k > n_samples {
            bail!("n_components must be <= min(n_features, n_samples)");
        }

        // Initialize W and H
        let seed = self.config.random_seed.unwrap_or(42);
        let w = if let Some(w_init) = w_init {
            if w_init.dim() != (n_features, k) {
                bail!("w_init must have shape ({}, {})", n_features, k);
            }
            w_init.clone()
        } else {
            random_nonneg_matrix(n_features, k, seed)
        };

        let h = random_nonneg_matrix(k, n_samples, seed + 1);

        // Upload matrices to GPU
        let gpu_v = self.backend.upload(v)?;
        let mut gpu_w = self.backend.upload(&w)?;
        let mut gpu_h = self.backend.upload(&h)?;

        // Pre-compute V transpose for efficient operations (reserved for future optimizations)
        let _gpu_v_t = self.backend.transpose(&gpu_v)?;

        let mut prev_error = f32::INFINITY;
        let mut converged = false;
        let mut iterations = 0;

        const EPSILON: f32 = 1e-10;

        for iter in 0..self.config.max_iterations {
            iterations = iter + 1;

            // Update H: H = H .* (W^T @ V) ./ (W^T @ W @ H + eps)
            let gpu_w_t = self.backend.transpose(&gpu_w)?;
            let numerator_h = self.backend.matmul(&gpu_w_t, &gpu_v)?;
            let wtw = self.backend.matmul(&gpu_w_t, &gpu_w)?;
            let denominator_h = self.backend.matmul(&wtw, &gpu_h)?;
            self.backend
                .nmf_update_inplace(&mut gpu_h, &numerator_h, &denominator_h, EPSILON)?;

            // Update W: W = W .* (V @ H^T) ./ (W @ H @ H^T + eps)
            let gpu_h_t = self.backend.transpose(&gpu_h)?;
            let numerator_w = self.backend.matmul(&gpu_v, &gpu_h_t)?;
            let wh = self.backend.matmul(&gpu_w, &gpu_h)?;
            let denominator_w = self.backend.matmul(&wh, &gpu_h_t)?;
            self.backend
                .nmf_update_inplace(&mut gpu_w, &numerator_w, &denominator_w, EPSILON)?;

            // Check convergence every 10 iterations
            if iter % 10 == 0 || iter == self.config.max_iterations - 1 {
                // Compute reconstruction: approx = W @ H
                let approx = self.backend.matmul(&gpu_w, &gpu_h)?;
                let error = self.backend.frobenius_diff(&gpu_v, &approx)?;

                debug!("Iteration {}: error = {:.6}", iter, error);

                if (prev_error - error).abs() < self.config.tolerance {
                    info!("Converged at iteration {}", iter);
                    converged = true;
                    break;
                }

                prev_error = error;
            }
        }

        // Download results
        self.backend.synchronize()?;
        let w_matrix = self.backend.download(&gpu_w)?;
        let h_matrix = self.backend.download(&gpu_h)?;

        info!(
            "GPU NMF completed: {} iterations, error = {:.6}, converged = {}",
            iterations, prev_error, converged
        );

        Ok(NmfOutput {
            w_matrix,
            h_matrix,
            iterations,
            reconstruction_error: prev_error,
            converged,
        })
    }

    /// Apply NMF to separate stains from optical density image
    ///
    /// # Arguments
    /// * `od_pixels` - Optical density values, shape (n_pixels, 3)
    /// * `stain_matrix` - Known stain vectors, shape (3, n_components)
    ///
    /// # Returns
    /// Concentration values for each stain, shape (n_pixels, n_components)
    pub fn unmix_stains(
        &self,
        od_pixels: &Array2<f32>,
        stain_matrix: &Array2<f32>,
    ) -> Result<Array2<f32>> {
        let (_n_pixels, n_channels) = od_pixels.dim();
        let _n_components = stain_matrix.ncols();

        if n_channels != 3 {
            bail!("Expected 3 channels (RGB), got {}", n_channels);
        }
        if stain_matrix.nrows() != 3 {
            bail!("Stain matrix must have 3 rows (RGB)");
        }

        // Transpose to (channels x pixels) for NMF
        let v = od_pixels.t().to_owned();

        // Use stain matrix as initial W
        let result = self.decompose(&v, Some(stain_matrix))?;

        // H is (n_components x n_pixels), transpose to (n_pixels x n_components)
        Ok(result.h_matrix.t().to_owned())
    }
}

/// Generate a random non-negative matrix
fn random_nonneg_matrix(rows: usize, cols: usize, seed: u64) -> Array2<f32> {
    use rand::Rng;
    use rand::SeedableRng;
    use rand_chacha::ChaCha8Rng;

    let mut rng = ChaCha8Rng::seed_from_u64(seed);
    Array2::from_shape_fn((rows, cols), |_| rng.r#gen::<f32>() * 0.1 + 0.01)
}

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

    fn test_config() -> NmfProcessingConfig {
        NmfProcessingConfig {
            n_components: 2,
            max_iterations: 50,
            tolerance: 1e-3,
            use_gpu: true,
            random_seed: Some(42),
        }
    }

    #[test]
    fn test_gpu_nmf_basic() {
        let config = test_config();
        let processor = GpuNmfProcessor::new(config).unwrap();

        // Create a simple test matrix
        let v = Array2::from_shape_vec(
            (3, 100),
            (0..300).map(|i| (i as f32 / 300.0) + 0.1).collect(),
        )
        .unwrap();

        let result = processor.decompose(&v, None).unwrap();

        assert_eq!(result.w_matrix.shape(), &[3, 2]);
        assert_eq!(result.h_matrix.shape(), &[2, 100]);
        assert!(result.reconstruction_error.is_finite());
    }

    #[test]
    fn test_gpu_nmf_with_init() {
        let config = test_config();
        let processor = GpuNmfProcessor::new(config).unwrap();

        // Create test matrix and initial W
        let v = Array2::from_shape_fn((3, 50), |(i, j)| ((i + j) as f32 / 10.0) + 0.1);
        let w_init = Array2::from_shape_fn((3, 2), |_| 0.5);

        let result = processor.decompose(&v, Some(&w_init)).unwrap();

        assert_eq!(result.w_matrix.shape(), &[3, 2]);
    }

    #[test]
    fn test_gpu_unmix_stains() {
        let config = test_config();
        let processor = GpuNmfProcessor::new(config).unwrap();

        // Create synthetic OD pixels
        let od_pixels =
            Array2::from_shape_fn((100, 3), |(i, j)| ((i * 3 + j) as f32 / 300.0) + 0.1);

        // Standard H&E stain matrix (simplified)
        let stain_matrix =
            Array2::from_shape_vec((3, 2), vec![0.65, 0.07, 0.70, 0.99, 0.29, 0.11]).unwrap();

        let concentrations = processor.unmix_stains(&od_pixels, &stain_matrix).unwrap();

        assert_eq!(concentrations.shape(), &[100, 2]);

        // Concentrations should be non-negative
        for &c in concentrations.iter() {
            assert!(c >= 0.0);
        }
    }

    #[test]
    fn test_backend_detection() {
        let config = test_config();
        let processor = GpuNmfProcessor::new(config).unwrap();

        // Should have some backend (at minimum CPU)
        let backend_type = processor.backend_type();
        println!("Using backend: {:?}", backend_type);

        // backend_type should be one of the valid types
        assert!(matches!(
            backend_type,
            BackendType::Cuda | BackendType::Metal | BackendType::Cpu
        ));
    }
}