// Simple test to verify GPU usage for basic operations
use rtx_nmf::{Device, Tensor};
use std::time::Instant;

fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("🔬 Simple GPU vs CPU Comparison");
    println!("==============================");

    // Test both CPU and GPU if available
    test_device_performance()?;

    Ok(())
}

fn test_device_performance() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n📱 Testing CPU performance (baseline)...");
    let cpu_device = Device::cpu();
    let cpu_time = test_matrix_operations(&cpu_device, "CPU")?;

    println!("\n📱 Testing GPU performance...");
    match Device::cuda(0) {
        Ok(gpu_device) => {
            let gpu_time = test_matrix_operations(&gpu_device, "GPU")?;

            println!("\n📊 Performance Comparison:");
            println!("   CPU time: {:.3}ms", cpu_time);
            println!("   GPU time: {:.3}ms", gpu_time);

            if gpu_time < cpu_time * 0.8 {
                println!(
                    "   ✅ GPU is faster than CPU by {:.1}x",
                    cpu_time / gpu_time
                );
            } else if gpu_time > cpu_time * 1.2 {
                println!("   🚨 WARNING: GPU is slower than CPU!");
                println!("   🔍 This indicates GPU operations are not working properly");
            } else {
                println!("   ⚠️  GPU and CPU performance are similar");
                println!("   🔍 This might indicate CPU fallback in GPU operations");
            }
        }
        Err(e) => {
            println!("   ❌ GPU not available: {:?}", e);
        }
    }

    Ok(())
}

fn test_matrix_operations(
    device: &Device,
    device_name: &str,
) -> Result<f32, Box<dyn std::error::Error>> {
    println!("   Creating test matrices on {}...", device_name);

    // Create smaller matrices to avoid issues
    let size = 512; // 512x512 matrices
    let data_a: Vec<f32> = (0..size * size).map(|i| (i % 100) as f32 + 1.0).collect();
    let data_b: Vec<f32> = (0..size * size)
        .map(|i| ((i * 2) % 100) as f32 + 1.0)
        .collect();

    let a = Tensor::from_data(data_a, [size, size], device)?;
    let b = Tensor::from_data(data_b, [size, size], device)?;

    println!(
        "   Testing {}x{} matrix multiplication on {}...",
        size, size, device_name
    );

    // Warm up (first call often slower)
    let _warmup = a.matmul(&b)?;

    // Actual timing
    let start = Instant::now();
    let _result = a.matmul(&b)?;
    let elapsed = start.elapsed();

    let time_ms = elapsed.as_secs_f32() * 1000.0;
    println!("   {} matmul time: {:.3}ms", device_name, time_ms);

    // Calculate theoretical performance
    let flops = 2.0 * size as f32 * size as f32 * size as f32; // 2*N^3 for NxN matmul
    let gflops = (flops / 1e9) / elapsed.as_secs_f32();
    println!("   {} achieved: {:.1} GFLOPS", device_name, gflops);

    if device_name == "GPU" {
        println!("   RTX 5090 theoretical: ~83,000 GFLOPS");
        println!("   Utilization: {:.3}%", (gflops / 83000.0) * 100.0);
    }

    Ok(time_ms)
}