// Test actual GPU performance and utilization
use rtx_nmf::{Device, NMFConfig, NMFDecomposer, NMFDemo, Tensor};
use std::time::Instant;

fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("🚀 GPU Performance Analysis for RTX-NMF");
    println!("=======================================");

    // First test GPU device availability
    println!("\n📱 Testing device availability...");

    match Device::cuda(0) {
        Ok(gpu_device) => {
            println!("✅ GPU Device available: {:?}", gpu_device);
            test_gpu_performance(&gpu_device)?;
        }
        Err(e) => {
            println!("⚠️  GPU not available: {:?}", e);
            println!("🔄 Falling back to CPU test...");
            test_cpu_performance()?;
        }
    }

    Ok(())
}

fn test_gpu_performance(device: &Device) -> Result<(), Box<dyn std::error::Error>> {
    println!("\n🔥 Testing GPU Performance");
    println!("==========================");

    // Test 1: Large tensor creation
    println!("\n⚡ Test 1: Large tensor creation (2000x2000)");
    let start = Instant::now();
    let large_matrix = create_large_matrix(2000, 2000, device)?;
    let creation_time = start.elapsed();
    println!(
        "   Creation time: {:.3}ms",
        creation_time.as_secs_f32() * 1000.0
    );
    println!(
        "   Matrix size: {} MB",
        (large_matrix.numel() * 4) / (1024 * 1024)
    );

    // Test 2: Matrix operations
    println!("\n⚡ Test 2: Matrix operations (1000x1000)");
    let a = create_large_matrix(1000, 1000, device)?;
    let b = create_large_matrix(1000, 1000, device)?;

    let start = Instant::now();
    let add_result = a.add(&b)?;
    let add_time = start.elapsed();

    let start = Instant::now();
    let mul_result = a.matmul(&b)?;
    let matmul_time = start.elapsed();

    println!("   Addition time: {:.3}ms", add_time.as_secs_f32() * 1000.0);
    println!(
        "   Matrix multiply time: {:.3}ms",
        matmul_time.as_secs_f32() * 1000.0
    );
    println!("   Results shape: {:?}", mul_result.shape().dims());

    // Test 3: NMF performance on medium image
    println!("\n⚡ Test 3: NMF algorithm performance");
    let mut demo = NMFDemo::new(device.clone());

    println!("   Running 128x128 image decomposition...");
    let start = Instant::now();
    let result = demo.run_image_demo(128, 128)?;
    let nmf_time = start.elapsed();

    println!(
        "   NMF total time: {:.3}ms",
        nmf_time.as_secs_f32() * 1000.0
    );
    println!(
        "   Algorithm time: {:.3}ms",
        result.computation_time * 1000.0
    );
    println!(
        "   Components: {}, Iterations: {}",
        result.components, result.iterations
    );
    println!(
        "   Reconstruction error: {:.6}",
        result.reconstruction_error
    );

    // Performance expectations for RTX 5090
    println!("\n📊 Performance Analysis:");
    println!("   Expected RTX 5090 1000x1000 matmul: ~1-5ms");
    println!(
        "   Actual matmul: {:.3}ms",
        matmul_time.as_secs_f32() * 1000.0
    );
    if matmul_time.as_secs_f32() * 1000.0 > 10.0 {
        println!("   🚨 WARNING: Performance significantly below RTX 5090 expectations");
        println!("   🔍 This suggests CPU execution or poor GPU utilization");
    } else {
        println!("   ✅ Performance within expected RTX 5090 range");
    }

    Ok(())
}

fn test_cpu_performance() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n💻 Testing CPU Performance (baseline)");
    println!("=====================================");

    let device = Device::cpu();

    // Test smaller matrices on CPU
    let a = create_large_matrix(500, 500, &device)?;
    let b = create_large_matrix(500, 500, &device)?;

    let start = Instant::now();
    let cpu_result = a.matmul(&b)?;
    let cpu_time = start.elapsed();

    println!(
        "   CPU 500x500 matmul: {:.3}ms",
        cpu_time.as_secs_f32() * 1000.0
    );
    println!("   Result shape: {:?}", cpu_result.shape().dims());

    Ok(())
}

fn create_large_matrix(
    rows: usize,
    cols: usize,
    device: &Device,
) -> Result<Tensor, Box<dyn std::error::Error>> {
    // Create matrix with some pattern to make it interesting
    let mut data = Vec::with_capacity(rows * cols);
    for i in 0..rows {
        for j in 0..cols {
            let value = (i as f32 * 0.1 + j as f32 * 0.2) % 10.0 + 1.0; // Positive values
            data.push(value);
        }
    }

    Ok(Tensor::from_data(data, [rows, cols], device)?)
}