//! GPU Memory Monitoring Tool
//!
//! Monitors VRAM usage during memory optimization tests
//! to verify 4GB GPU compatibility.

use candle_core::{Device, Tensor};
use ml::memory_optimization::{
    PrecisionConverter, PrecisionType, QuantizationConfig, QuantizationType, Quantizer,
};
use std::process::Command;
use std::thread;
use std::time::{Duration, Instant};

fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== GPU Memory Monitor for 4GB RTX 3050 Ti ===\n");

    // Check initial GPU memory
    print_gpu_memory("Initial State")?;

    let device = Device::cuda_if_available(0)?;
    println!("Device: {:?}\n", device);

    // Test 1: Baseline memory usage
    test_baseline_memory(&device)?;

    // Test 2: Large tensor allocation
    test_large_tensor_memory(&device)?;

    // Test 3: Multiple models
    test_multiple_models(&device)?;

    // Test 4: Memory optimization impact
    test_optimization_impact(&device)?;

    println!("\n=== GPU Memory Monitoring Complete ===");
    Ok(())
}

fn print_gpu_memory(label: &str) -> Result<(), Box<dyn std::error::Error>> {
    println!("--- {} ---", label);

    // Run nvidia-smi to get GPU memory info
    let output = Command::new("nvidia-smi")
        .args(&[
            "--query-gpu=memory.used,memory.free,memory.total",
            "--format=csv,noheader,nounits",
        ])
        .output()?;

    if output.status.success() {
        let result = String::from_utf8_lossy(&output.stdout);
        let parts: Vec<&str> = result.trim().split(", ").collect();

        if parts.len() == 3 {
            let used: f64 = parts[0].parse().unwrap_or(0.0);
            let free: f64 = parts[1].parse().unwrap_or(0.0);
            let total: f64 = parts[2].parse().unwrap_or(0.0);

            println!("GPU Memory:");
            println!("  Used:  {:.0} MB", used);
            println!("  Free:  {:.0} MB", free);
            println!("  Total: {:.0} MB", total);
            println!("  Usage: {:.1}%", (used / total) * 100.0);
        }
    } else {
        println!("nvidia-smi not available");
    }

    println!();
    Ok(())
}

fn test_baseline_memory(device: &Device) -> Result<(), Box<dyn std::error::Error>> {
    println!("Test 1: Baseline Memory Usage");
    println!("-------------------------------");

    let start = Instant::now();

    // Create a small tensor
    let tensor = Tensor::randn(0.0f32, 1.0f32, (100, 100), device)?;
    let size_mb = (tensor.dims().iter().product::<usize>() * 4) as f64 / 1_048_576.0;

    println!("Created tensor: {:?}, size: {:.2} MB", tensor.dims(), size_mb);

    thread::sleep(Duration::from_millis(500));
    print_gpu_memory("After Small Tensor")?;

    drop(tensor);
    thread::sleep(Duration::from_millis(500));

    let elapsed = start.elapsed();
    println!("✓ Baseline test complete ({:.2}ms)\n", elapsed.as_secs_f64() * 1000.0);

    Ok(())
}

fn test_large_tensor_memory(device: &Device) -> Result<(), Box<dyn std::error::Error>> {
    println!("Test 2: Large Tensor Memory Usage");
    println!("-----------------------------------");

    let start = Instant::now();

    // Allocate progressively larger tensors
    let sizes = vec![
        (256, 256),
        (512, 512),
        (1024, 1024),
        (2048, 2048),
    ];

    for (h, w) in sizes {
        let tensor = Tensor::randn(0.0f32, 1.0f32, (h, w), device)?;
        let size_mb = (tensor.dims().iter().product::<usize>() * 4) as f64 / 1_048_576.0;

        println!("Tensor [{}, {}]: {:.2} MB", h, w, size_mb);

        thread::sleep(Duration::from_millis(200));

        drop(tensor);
    }

    thread::sleep(Duration::from_millis(500));
    print_gpu_memory("After Large Tensors")?;

    let elapsed = start.elapsed();
    println!("✓ Large tensor test complete ({:.2}ms)\n", elapsed.as_secs_f64() * 1000.0);

    Ok(())
}

fn test_multiple_models(device: &Device) -> Result<(), Box<dyn std::error::Error>> {
    println!("Test 3: Multiple Model Simulation");
    println!("-----------------------------------");

    let start = Instant::now();

    // Simulate multiple models loaded simultaneously
    let model_configs = vec![
        ("DQN", 256, 256),
        ("PPO", 512, 256),
        ("MAMBA-2", 1024, 512),
    ];

    let mut tensors = Vec::new();

    for (name, h, w) in model_configs {
        let tensor = Tensor::randn(0.0f32, 1.0f32, (h, w), device)?;
        let size_mb = (tensor.dims().iter().product::<usize>() * 4) as f64 / 1_048_576.0;

        println!("{} model: [{}, {}] = {:.2} MB", name, h, w, size_mb);

        tensors.push(tensor);
    }

    thread::sleep(Duration::from_millis(500));
    print_gpu_memory("With Multiple Models")?;

    drop(tensors);
    thread::sleep(Duration::from_millis(500));

    let elapsed = start.elapsed();
    println!("✓ Multiple models test complete ({:.2}ms)\n", elapsed.as_secs_f64() * 1000.0);

    Ok(())
}

fn test_optimization_impact(device: &Device) -> Result<(), Box<dyn std::error::Error>> {
    println!("Test 4: Memory Optimization Impact");
    println!("------------------------------------");

    let start = Instant::now();

    // Test baseline F32
    println!("\n[Phase 1: Baseline F32]");
    let tensor_f32 = Tensor::randn(0.0f32, 1.0f32, (1024, 1024), device)?;
    let size_f32 = (tensor_f32.dims().iter().product::<usize>() * 4) as f64 / 1_048_576.0;
    println!("F32 tensor size: {:.2} MB", size_f32);

    thread::sleep(Duration::from_millis(500));
    print_gpu_memory("F32 Baseline")?;

    // Test FP16
    println!("[Phase 2: FP16 Conversion]");
    let mut converter = PrecisionConverter::new(PrecisionType::Float16, device.clone());
    let tensor_f16 = converter.to_float16(&tensor_f32)?;
    let size_f16 = (tensor_f16.dims().iter().product::<usize>() * 2) as f64 / 1_048_576.0;
    println!("F16 tensor size: {:.2} MB (saved {:.2} MB)", size_f16, size_f32 - size_f16);

    thread::sleep(Duration::from_millis(500));
    print_gpu_memory("After FP16")?;

    // Test INT8 quantization
    println!("[Phase 3: INT8 Quantization]");
    let tensor_for_quant = converter.to_float32(&tensor_f16)?;

    let quant_config = QuantizationConfig {
        quant_type: QuantizationType::Int8,
        symmetric: true,
        per_channel: true,
        calibration_samples: Some(1000),
    };

    let mut quantizer = Quantizer::new(quant_config, device.clone());
    let quantized = quantizer.quantize_tensor(&tensor_for_quant, "test_model")?;

    let size_quant = quantized.memory_bytes() as f64 / 1_048_576.0;
    println!("INT8 tensor size: {:.2} MB (saved {:.2} MB from baseline)", size_quant, size_f32 - size_quant);

    thread::sleep(Duration::from_millis(500));
    print_gpu_memory("After INT8 Quantization")?;

    // Summary
    println!("\n--- Optimization Summary ---");
    println!("Baseline (F32):    {:.2} MB (100.0%)", size_f32);
    println!("FP16:              {:.2} MB ({:.1}%)", size_f16, (size_f16 / size_f32) * 100.0);
    println!("INT8:              {:.2} MB ({:.1}%)", size_quant, (size_quant / size_f32) * 100.0);
    println!("Total Savings:     {:.2} MB ({:.1}%)", size_f32 - size_quant, ((size_f32 - size_quant) / size_f32) * 100.0);

    let elapsed = start.elapsed();
    println!("\n✓ Optimization impact test complete ({:.2}ms)\n", elapsed.as_secs_f64() * 1000.0);

    Ok(())
}