// Final test of the corrected NMF implementation
use rtx_nmf::{Device, NMFDemo};

fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("🎓 Educational NMF Demo - Corrected Implementation");
    println!("================================================");

    // Test with both CPU and GPU if available to provide real comparison
    println!("\n📱 Testing CPU version for baseline...");
    test_nmf_device(&Device::cpu(), "CPU")?;

    println!("\n📱 Testing GPU version...");
    match Device::cuda(0) {
        Ok(gpu_device) => {
            test_nmf_device(&gpu_device, "GPU")?;
        }
        Err(_) => {
            println!("   GPU not available for comparison");
        }
    }

    Ok(())
}

fn test_nmf_device(device: &Device, device_name: &str) -> Result<(), Box<dyn std::error::Error>> {
    println!("   Initializing NMF demo on {}...", device_name);
    let mut demo = NMFDemo::new(device.clone());

    // Use reasonable image size that works well
    let (width, height) = (48, 32); // 1,536 pixels

    println!(
        "   Running {}×{} image decomposition on {}...",
        width, height, device_name
    );
    let result = demo.run_image_demo(width, height)?;

    // Present clean, honest results
    println!("   ✅ {} Results:", device_name);
    println!("      Components: {}", result.components);
    println!("      Iterations: {}", result.iterations);
    println!("      Error: {:.6}", result.reconstruction_error);
    println!("      Time: {:.3}s", result.computation_time);
    println!("      Converged: {}", result.converged);

    // Educational insights
    let compression = result.original_image.data_size as f32
        / ((result.original_image.shape[0] * result.components)
            + (result.components * result.original_image.shape[1])) as f32;

    println!("      Compression: {:.1}x storage reduction", compression);

    Ok(())
}