//! Debug GPU NMF implementation to isolate division by zero error

use rtx_nmf::{NMFConfig, NMFDecomposer};
use rtx_tensor::{Device, Tensor};

fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("🔧 Debugging GPU NMF Implementation");
    println!("===================================");

    // Test GPU device
    let device = Device::cuda(0).unwrap_or_else(|e| {
        println!("⚠️ GPU not available: {:?}, using CPU", e);
        Device::cpu()
    });

    println!("🔥 Using device: {:?}", device);

    // Create simple test matrix (non-negative)
    println!("\n🧪 Creating simple test matrix...");
    let test_data = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 2.0, 4.0, 6.0];

    let matrix = Tensor::from_vec(test_data, &[4, 3], &device)?;
    println!("   Matrix shape: {:?}", matrix.shape());
    println!("   Matrix device: {:?}", matrix.device());

    // Test tensor data access
    println!("\n🔍 Testing tensor data access...");
    let data_check = matrix.to_vec()?;
    println!("   Data length: {} (expected: {})", data_check.len(), 4 * 3);
    println!(
        "   Data range: [{:.3}, {:.3}]",
        data_check.iter().fold(f32::INFINITY, |a, &b| a.min(b)),
        data_check.iter().fold(f32::NEG_INFINITY, |a, &b| a.max(b))
    );

    // Test basic tensor operations
    println!("\n⚗️ Testing basic tensor operations...");

    // Test scalar addition
    let plus_epsilon = matrix.add_scalar(1e-6)?;
    println!("   Scalar addition: ✅");

    // Test matrix multiplication
    let transpose = matrix.transpose(0, 1)?;
    println!("   Transpose: ✅ shape {:?}", transpose.shape());

    let matmul_test = transpose.matmul(&matrix)?;
    println!(
        "   Matrix multiplication: ✅ shape {:?}",
        matmul_test.shape()
    );

    // Test division (this is where the error likely occurs)
    println!("\n🧮 Testing division operation...");
    let ones = Tensor::from_vec(vec![1.0; 9], &[3, 3], &device)?;
    let small_values = ones.add_scalar(1e-6)?;

    match ones.div(&small_values) {
        Ok(result) => {
            println!("   Simple division: ✅");
            let result_data = result.to_vec()?;
            println!(
                "   Division result range: [{:.3}, {:.3}]",
                result_data.iter().fold(f32::INFINITY, |a, &b| a.min(b)),
                result_data.iter().fold(f32::NEG_INFINITY, |a, &b| a.max(b))
            );
        }
        Err(e) => {
            println!("   ❌ Division failed: {:?}", e);
            return Err(e.into());
        }
    }

    // Now test NMF with small matrices
    println!("\n🧮 Testing NMF algorithm...");
    let config = NMFConfig::new()
        .with_components(2) // Reduced components
        .with_max_iterations(5) // Reduced iterations
        .with_epsilon(1e-4); // Larger epsilon

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

    println!("   NMF Config: 2 components, 5 iterations, epsilon=1e-4");

    match nmf.fit_transform_detailed(&matrix) {
        Ok(result) => {
            println!("   ✅ NMF succeeded!");
            println!("   W shape: {:?}", result.w.shape());
            println!("   H shape: {:?}", result.h.shape());
            println!("   Iterations: {}", result.iterations);
            println!("   Error: {:.6}", result.reconstruction_error);
            println!("   Time: {:.3}s", result.computation_time);
        }
        Err(e) => {
            println!("   ❌ NMF failed: {:?}", e);

            // Try with CPU device instead
            println!("\n🔄 Retrying with CPU device...");
            let cpu_device = Device::cpu();
            let cpu_matrix = Tensor::from_vec(
                vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 2.0, 4.0, 6.0],
                &[4, 3],
                &cpu_device,
            )?;

            let cpu_config = NMFConfig::new()
                .with_components(2)
                .with_max_iterations(5)
                .with_epsilon(1e-4);

            let mut cpu_nmf = NMFDecomposer::new(cpu_config);

            match cpu_nmf.fit_transform_detailed(&cpu_matrix) {
                Ok(cpu_result) => {
                    println!("   ✅ CPU NMF succeeded!");
                    println!("   W shape: {:?}", cpu_result.w.shape());
                    println!("   H shape: {:?}", cpu_result.h.shape());
                    println!("   GPU vs CPU: GPU failed, CPU worked");
                }
                Err(cpu_e) => {
                    println!("   ❌ CPU NMF also failed: {:?}", cpu_e);
                }
            }
        }
    }

    Ok(())
}