// Debug GPU image generation issue
use rtx_nmf::demo::ImageProcessor;
use rtx_tensor::{Device, Tensor};

fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("🔍 Debugging GPU Image Generation Issue");
    println!("======================================");

    // Compare CPU vs GPU image generation
    println!("\n📱 Testing CPU image generation...");
    let cpu_device = Device::cpu();
    let cpu_processor = ImageProcessor::new(cpu_device);
    let cpu_image = cpu_processor.create_mock_image(16, 12)?;
    let cpu_info = cpu_processor.tensor_to_image_info(&cpu_image)?;

    println!("   CPU Image:");
    println!("     Shape: {:?}", cpu_info.shape);
    println!(
        "     Range: [{:.3}, {:.3}]",
        cpu_info.min_value, cpu_info.max_value
    );
    println!("     Mean: {:.3}", cpu_info.mean_value);

    // Test raw data
    let cpu_data = cpu_image.to_cpu()?;
    let non_zero_count = cpu_data.iter().filter(|&&x| x > 0.0).count();
    println!(
        "     Non-zero elements: {}/{}",
        non_zero_count,
        cpu_data.len()
    );

    println!("\n📱 Testing GPU image generation...");
    match Device::cuda(0) {
        Ok(gpu_device) => {
            let gpu_processor = ImageProcessor::new(gpu_device);
            let gpu_image = gpu_processor.create_mock_image(16, 12)?;
            let gpu_info = gpu_processor.tensor_to_image_info(&gpu_image)?;

            println!("   GPU Image:");
            println!("     Shape: {:?}", gpu_info.shape);
            println!(
                "     Range: [{:.3}, {:.3}]",
                gpu_info.min_value, gpu_info.max_value
            );
            println!("     Mean: {:.3}", gpu_info.mean_value);

            // Test raw data
            let gpu_data = gpu_image.to_cpu()?; // Transfer to CPU for inspection
            let gpu_non_zero_count = gpu_data.iter().filter(|&&x| x > 0.0).count();
            println!(
                "     Non-zero elements: {}/{}",
                gpu_non_zero_count,
                gpu_data.len()
            );

            // Check if data is identical
            if cpu_data.len() == gpu_data.len() {
                let differences: usize = cpu_data
                    .iter()
                    .zip(gpu_data.iter())
                    .filter(|&(a, b)| (a - b).abs() > 1e-6)
                    .count();
                println!(
                    "     Differences from CPU: {}/{}",
                    differences,
                    cpu_data.len()
                );

                if gpu_info.max_value == 0.0 {
                    println!("   🚨 PROBLEM: GPU image generation produces all zeros!");
                    println!("   🔍 This suggests the image generation logic fails on GPU");
                } else if differences == 0 {
                    println!("   ✅ GPU and CPU produce identical results");
                } else {
                    println!("   ⚠️  GPU and CPU produce different results");
                }
            }
        }
        Err(e) => {
            println!("   ❌ GPU not available: {:?}", e);
        }
    }

    // Test simple tensor creation on both devices
    println!("\n🧮 Testing basic tensor operations...");
    test_basic_tensor_ops()?;

    Ok(())
}

fn test_basic_tensor_ops() -> Result<(), Box<dyn std::error::Error>> {
    println!("   Creating simple tensor on CPU...");
    let cpu_device = Device::cpu();
    let cpu_tensor = Tensor::from_data(vec![1.0, 2.0, 3.0, 4.0], [2, 2], &cpu_device)?;
    let cpu_data = cpu_tensor.to_cpu()?;
    println!("     CPU tensor: {:?}", cpu_data);

    match Device::cuda(0) {
        Ok(gpu_device) => {
            println!("   Creating simple tensor on GPU...");
            let gpu_tensor = Tensor::from_data(vec![1.0, 2.0, 3.0, 4.0], [2, 2], &gpu_device)?;
            let gpu_data = gpu_tensor.to_cpu()?; // Transfer back
            println!("     GPU tensor (transferred to CPU): {:?}", gpu_data);

            if cpu_data == gpu_data {
                println!("   ✅ Basic tensor operations work correctly on GPU");
            } else {
                println!("   🚨 PROBLEM: GPU tensor operations produce different results!");
            }
        }
        Err(_) => {
            println!("   GPU not available for tensor test");
        }
    }

    Ok(())
}