//! Comprehensive tests for rtx-nmf (Non-negative Matrix Factorization)
//! Tests both the GPU-accelerated NMFDecomposer and CPU-based HonestNMF

use rtx_nmf::*;
use rtx_tensor::{Device, Tensor};

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

    #[test]
    fn test_nmf_config_creation_with_defaults() {
        // Test: NMFConfig should initialize with default parameters
        let config = NMFConfig::new();

        assert_eq!(config.components(), 10);
        assert_eq!(config.max_iterations(), 100);
        assert_eq!(config.tolerance(), 1e-4);
    }

    #[test]
    fn test_nmf_config_with_custom_parameters() {
        // Test: NMFConfig should accept custom parameters
        let config = NMFConfig::new()
            .with_components(5)
            .with_max_iterations(200)
            .with_tolerance(1e-6)
            .with_epsilon(1e-8)
            .with_random_seed(42);

        assert_eq!(config.components(), 5);
        assert_eq!(config.max_iterations(), 200);
        assert_eq!(config.tolerance(), 1e-6);
        assert_eq!(config.epsilon(), 1e-8);
    }

    #[test]
    fn test_nmf_config_validation() {
        // Test: Config validation should catch invalid parameters
        let valid_config = NMFConfig::new().with_components(5);
        assert!(valid_config.validate().is_ok());

        let invalid_config = NMFConfig::new().with_components(0);
        assert!(invalid_config.validate().is_err());
    }
}

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

    #[test]
    fn test_basic_nmf_decomposition() -> Result<()> {
        // Test: NMFDecomposer should decompose a matrix into W and H
        let device = Device::Cuda(0);

        // Create a simple non-negative matrix
        let data = vec![
            1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0,
        ];

        let x = Tensor::from_data(data, [4, 3], &device)?;

        let config = NMFConfig::new().with_components(2).with_max_iterations(100);
        let mut nmf = NMFDecomposer::new(config);

        let (w, h) = nmf.fit_transform(&x)?;

        // Check dimensions
        assert_eq!(w.shape().dims(), &[4, 2]); // n_samples x n_components
        assert_eq!(h.shape().dims(), &[2, 3]); // n_components x n_features

        // Verify non-negativity
        let w_data = w.to_cpu()?;
        let h_data = h.to_cpu()?;

        for val in w_data {
            assert!(val >= 0.0, "W matrix should be non-negative");
        }

        for val in h_data {
            assert!(val >= 0.0, "H matrix should be non-negative");
        }

        Ok(())
    }

    #[test]
    fn test_nmf_reconstruction() -> Result<()> {
        // Test: Reconstruction should approximate original matrix
        let device = Device::Cuda(0);

        // Create synthetic data
        let data: Vec<f32> = (0..50).map(|i| (i as f32 + 1.0) * 0.5).collect();
        let x = Tensor::from_data(data, [10, 5], &device)?;

        let config = NMFConfig::new().with_components(3).with_max_iterations(100);
        let mut nmf = NMFDecomposer::new(config);

        let result = nmf.fit_transform_detailed(&x)?;

        // Calculate reconstruction
        let reconstruction = result.w.matmul(&result.h)?;

        // Calculate difference
        let diff = x.sub(&reconstruction)?;
        let error = diff.matrix_norm("fro")?;

        // Error should be reasonably small
        assert!(
            error < 10.0,
            "Reconstruction error should be small, got {}",
            error
        );
        assert!(result.reconstruction_error >= 0.0);

        Ok(())
    }

    #[test]
    fn test_nmf_convergence_info() -> Result<()> {
        // Test: NMF should provide convergence information
        let device = Device::Cuda(0);

        let data: Vec<f32> = (0..200).map(|i| (i as f32).abs()).collect();
        let x = Tensor::from_data(data, [20, 10], &device)?;

        let config = NMFConfig::new()
            .with_components(5)
            .with_max_iterations(50)
            .with_tolerance(1e-4);
        let mut nmf = NMFDecomposer::new(config);

        let result = nmf.fit_transform_detailed(&x)?;

        // Check result fields
        assert!(result.iterations > 0 && result.iterations <= 50);
        assert!(result.reconstruction_error >= 0.0);
        assert!(result.computation_time > 0.0);

        Ok(())
    }
}

#[cfg(test)]
mod honest_nmf_tests {
    use super::*;
    use nalgebra::DMatrix;

    #[test]
    fn test_honest_nmf_basic() -> Result<()> {
        // Test: HonestNMF should work with CPU matrices
        let config = HonestNMFConfig::new()
            .with_components(3)
            .with_max_iterations(100)
            .with_tolerance(1e-4);

        let nmf = HonestNMF::new(config)?;

        // Create test matrix
        let data = vec![
            1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0, 13.0, 14.0, 15.0, 16.0,
            17.0, 18.0, 19.0, 20.0,
        ];
        let matrix = DMatrix::from_row_slice(5, 4, &data);

        let result = nmf.fit_transform(&matrix)?;

        // Check dimensions
        assert_eq!(result.w.nrows(), 5);
        assert_eq!(result.w.ncols(), 3);
        assert_eq!(result.h.nrows(), 3);
        assert_eq!(result.h.ncols(), 4);

        // Verify non-negativity
        for val in result.w.iter() {
            assert!(*val >= 0.0, "W should be non-negative");
        }
        for val in result.h.iter() {
            assert!(*val >= 0.0, "H should be non-negative");
        }

        Ok(())
    }

    #[test]
    fn test_honest_nmf_reconstruction() -> Result<()> {
        // Test: HonestNMF reconstruction
        let config = HonestNMFConfig::new()
            .with_components(2)
            .with_max_iterations(50);

        let nmf = HonestNMF::new(config)?;

        let data = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0];
        let matrix = DMatrix::from_row_slice(3, 3, &data);

        let result = nmf.fit_transform(&matrix)?;
        let reconstruction = result.reconstruct();

        // Check dimensions match
        assert_eq!(reconstruction.nrows(), matrix.nrows());
        assert_eq!(reconstruction.ncols(), matrix.ncols());

        // Error should be finite and non-negative
        assert!(result.reconstruction_error.is_finite());
        assert!(result.reconstruction_error >= 0.0);

        Ok(())
    }

    #[test]
    fn test_honest_nmf_convergence() -> Result<()> {
        // Test: HonestNMF convergence tracking
        let config = HonestNMFConfig::new()
            .with_components(4)
            .with_max_iterations(100)
            .with_tolerance(1e-4)
            .with_random_seed(42);

        let nmf = HonestNMF::new(config)?;

        let data: Vec<f32> = (0..100).map(|i| (i as f32) * 0.5).collect();
        let matrix = DMatrix::from_row_slice(10, 10, &data);

        let result = nmf.fit_transform(&matrix)?;

        assert!(result.iterations > 0);
        assert!(result.iterations <= 100);
        assert!(result.computation_time > 0.0);

        Ok(())
    }
}

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

    #[test]
    fn test_random_initialization() -> Result<()> {
        // Test: Random initialization strategy
        let device = Device::Cuda(0);

        let config = NMFConfig::new()
            .with_components(3)
            .with_initialization("random")
            .with_random_seed(42);

        let mut nmf = NMFDecomposer::new(config);

        let data: Vec<f32> = (0..60).map(|i| (i as f32) + 1.0).collect();
        let x = Tensor::from_data(data, [10, 6], &device)?;

        let (w, h) = nmf.fit_transform(&x)?;

        // Should produce valid matrices
        assert_eq!(w.shape().dims(), &[10, 3]);
        assert_eq!(h.shape().dims(), &[3, 6]);

        Ok(())
    }
}

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

    #[test]
    fn test_gpu_vs_cpu_mode() -> Result<()> {
        // Test: Both GPU and CPU modes should work
        let device = Device::Cuda(0);

        let data: Vec<f32> = (0..30).map(|i| (i as f32) + 0.5).collect();
        let x = Tensor::from_data(data, [6, 5], &device)?;

        // GPU mode (default)
        let config_gpu = NMFConfig::new().with_components(2).with_max_iterations(50);
        let mut nmf_gpu = NMFDecomposer::new(config_gpu);
        let result_gpu = nmf_gpu.fit_transform_detailed(&x)?;

        // CPU mode
        let config_cpu = NMFConfig::new()
            .with_components(2)
            .with_max_iterations(50)
            .without_gpu();
        let mut nmf_cpu = NMFDecomposer::new(config_cpu);
        let result_cpu = nmf_cpu.fit_transform_detailed(&x)?;

        // Both should produce valid results
        assert!(result_gpu.reconstruction_error >= 0.0);
        assert!(result_cpu.reconstruction_error >= 0.0);

        Ok(())
    }
}

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

    #[test]
    fn test_end_to_end_nmf_pipeline() -> Result<()> {
        // Test: Complete NMF pipeline
        let device = Device::Cuda(0);

        let n_samples = 50;
        let n_features = 20;
        let n_components = 5;

        // Generate synthetic data
        let data: Vec<f32> = (0..(n_samples * n_features))
            .map(|i| ((i % 17) as f32) * 0.7 + 0.1)
            .collect();

        let x = Tensor::from_data(data, [n_samples, n_features], &device)?;

        // Build and train NMF model
        let config = NMFConfig::new()
            .with_components(n_components)
            .with_max_iterations(100)
            .with_tolerance(1e-4)
            .with_random_seed(42);

        let mut nmf = NMFDecomposer::new(config);

        // Fit the model
        let result = nmf.fit_transform_detailed(&x)?;

        // Verify dimensions
        assert_eq!(result.w.shape().dims(), &[n_samples, n_components]);
        assert_eq!(result.h.shape().dims(), &[n_components, n_features]);

        // Reconstruct and check quality
        let reconstruction = result.reconstruct()?;
        assert_eq!(reconstruction.shape().dims(), &[n_samples, n_features]);

        // Check convergence
        assert!(result.iterations > 0);
        assert!(result.reconstruction_error >= 0.0);
        assert!(result.computation_time > 0.0);

        Ok(())
    }

    #[test]
    fn test_nmf_with_various_sizes() -> Result<()> {
        // Test: NMF should handle various matrix sizes
        let device = Device::Cuda(0);

        let test_cases = vec![
            (5, 3, 2),   // small
            (10, 8, 3),  // medium
            (20, 15, 5), // larger
        ];

        for (m, n, k) in test_cases {
            let data: Vec<f32> = (0..(m * n)).map(|i| (i as f32) * 0.1 + 0.5).collect();
            let x = Tensor::from_data(data, [m, n], &device)?;

            let config = NMFConfig::new().with_components(k).with_max_iterations(50);
            let mut nmf = NMFDecomposer::new(config);

            let (w, h) = nmf.fit_transform(&x)?;

            assert_eq!(w.shape().dims(), &[m, k]);
            assert_eq!(h.shape().dims(), &[k, n]);
        }

        Ok(())
    }
}

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

    #[test]
    fn test_invalid_components() {
        // Test: Should reject invalid number of components
        let config = NMFConfig::new().with_components(0);
        assert!(config.validate().is_err());
    }

    #[test]
    fn test_components_larger_than_matrix() -> Result<()> {
        // Test: Should handle components >= min(m, n)
        let device = Device::Cuda(0);

        let data: Vec<f32> = (0..12).map(|i| (i as f32) + 1.0).collect();
        let x = Tensor::from_data(data, [4, 3], &device)?;

        // Components >= min(4, 3) = 3 should fail
        let config = NMFConfig::new().with_components(3).with_max_iterations(10);
        let mut nmf = NMFDecomposer::new(config);

        let result = nmf.fit_transform(&x);
        assert!(result.is_err());

        Ok(())
    }
}

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

    #[test]
    fn test_nmf_demo_creation() {
        // Test: NMFDemo should be creatable
        let device = Device::Cuda(0);
        let _demo = NMFDemo::new(device);
        // Just verify it compiles and constructs
    }

    #[test]
    fn test_honest_nmf_demo() -> Result<()> {
        // Test: HonestNMFDemo basic functionality
        let _demo = HonestNMFDemo::new()?;
        // Verify construction
        Ok(())
    }
}