foxhunt/ml/examples/gpu_memory_benchmark.rs

/// GPU Memory Benchmarking Tool for RTX 3050 Ti (4GB VRAM)
///
/// This tool directly measures VRAM usage using nvidia-smi for accurate
/// GPU memory profiling of all ML models. Unlike system memory profiling,
/// this measures actual GPU VRAM allocations.
///
/// Run with: cargo run -p ml --example gpu_memory_benchmark --release --features cuda
///
/// Output: GPU_MEMORY_PROFILE_REPORT.md with VRAM budgets and batch size limits
use anyhow::{Context, Result};
use candle_core::{DType, Device, Tensor};
use chrono::Utc;
use std::collections::HashMap;
use std::fs::File;
use std::io::Write;
use std::path::PathBuf;
use std::process::Command;

// Expected memory ranges (MB) for each model
const DQN_RANGE_MB: (f64, f64) = (50.0, 150.0);
const PPO_RANGE_MB: (f64, f64) = (50.0, 200.0);
const MAMBA2_RANGE_MB: (f64, f64) = (150.0, 500.0);
const TFT_RANGE_MB: (f64, f64) = (1500.0, 2500.0);
const LIQUID_RANGE_MB: (f64, f64) = (100.0, 300.0);

/// GPU memory snapshot from nvidia-smi
#[derive(Debug, Clone)]
struct GpuMemorySnapshot {
    timestamp: chrono::DateTime<Utc>,
    total_mb: f64,
    free_mb: f64,
    used_mb: f64,
    utilization_percent: f64,
}

/// Model memory profile with GPU-specific metrics
#[derive(Debug, Clone)]
struct GpuModelProfile {
    model_name: String,
    base_vram_mb: f64,
    peak_vram_mb: f64,
    batch_size_tests: Vec<BatchTest>,
    max_safe_batch_size: u32,
    training_batch_size: u32,
    inference_batch_size: u32,
    parameter_count: usize,
    status: ProfileStatus,
}

#[derive(Debug, Clone)]
struct BatchTest {
    batch_size: u32,
    vram_used_mb: f64,
    success: bool,
    error_message: Option<String>,
}

#[derive(Debug, Clone, Copy, PartialEq)]
enum ProfileStatus {
    Safe,     // < 80% VRAM usage
    Tight,    // 80-90% VRAM usage
    Critical, // > 90% VRAM usage
    Failed,   // OOM or error
}

impl ProfileStatus {
    fn from_usage(used_mb: f64, total_mb: f64) -> Self {
        let usage_percent = (used_mb / total_mb) * 100.0;
        if usage_percent < 80.0 {
            ProfileStatus::Safe
        } else if usage_percent < 90.0 {
            ProfileStatus::Tight
        } else {
            ProfileStatus::Critical
        }
    }

    fn emoji(&self) -> &'static str {
        match self {
            ProfileStatus::Safe => "✅",
            ProfileStatus::Tight => "⚠️",
            ProfileStatus::Critical => "🔴",
            ProfileStatus::Failed => "❌",
        }
    }
}

/// Query GPU memory using nvidia-smi
fn query_gpu_memory() -> Result<GpuMemorySnapshot> {
    let output = Command::new("nvidia-smi")
        .args(&[
            "--query-gpu=memory.total,memory.free,memory.used,utilization.gpu",
            "--format=csv,noheader,nounits",
        ])
        .output()
        .context("Failed to execute nvidia-smi")?;

    let stdout = String::from_utf8(output.stdout).context("Failed to parse nvidia-smi output")?;

    let parts: Vec<&str> = stdout.trim().split(',').collect();

    if parts.len() < 4 {
        anyhow::bail!("Unexpected nvidia-smi output format: {}", stdout);
    }

    Ok(GpuMemorySnapshot {
        timestamp: Utc::now(),
        total_mb: parts[0].trim().parse::<f64>()?,
        free_mb: parts[1].trim().parse::<f64>()?,
        used_mb: parts[2].trim().parse::<f64>()?,
        utilization_percent: parts[3].trim().parse::<f64>()?,
    })
}

/// Profile DQN model VRAM usage
fn profile_dqn_vram(device: &Device, gpu_total_mb: f64) -> Result<GpuModelProfile> {
    println!("📊 Profiling DQN VRAM usage...");

    let input_dim = 64;
    let hidden_dim = 256;
    let action_dim = 5;

    // Measure baseline
    let baseline = query_gpu_memory()?;
    let base_vram_mb = baseline.used_mb;

    // Create model layers
    let layer1_w = Tensor::randn(0f32, 1.0, (hidden_dim, input_dim), device)?;
    let layer1_b = Tensor::zeros((hidden_dim,), DType::F32, device)?;
    let layer2_w = Tensor::randn(0f32, 1.0, (hidden_dim, hidden_dim), device)?;
    let layer2_b = Tensor::zeros((hidden_dim,), DType::F32, device)?;
    let layer3_w = Tensor::randn(0f32, 1.0, (action_dim, hidden_dim), device)?;
    let layer3_b = Tensor::zeros((action_dim,), DType::F32, device)?;

    std::thread::sleep(std::time::Duration::from_millis(100));
    let after_model = query_gpu_memory()?;
    println!(
        "  Model loaded: {:.1} MB VRAM",
        after_model.used_mb - base_vram_mb
    );

    // Test batch sizes
    let mut batch_tests = Vec::new();
    let batch_sizes = vec![1, 8, 16, 32, 64, 128, 256, 512];

    for batch_size in batch_sizes {
        print!("  Testing batch size {}: ", batch_size);

        let result = (|| -> Result<f64> {
            let input = Tensor::randn(0f32, 1.0, (batch_size, input_dim), device)?;
            let hidden1 = input.matmul(&layer1_w.t()?)?.broadcast_add(&layer1_b)?;
            let relu1 = hidden1.relu()?;
            let hidden2 = relu1.matmul(&layer2_w.t()?)?.broadcast_add(&layer2_b)?;
            let relu2 = hidden2.relu()?;
            let _output = relu2.matmul(&layer3_w.t()?)?.broadcast_add(&layer3_b)?;

            std::thread::sleep(std::time::Duration::from_millis(100));
            let mem = query_gpu_memory()?;
            Ok(mem.used_mb)
        })();

        match result {
            Ok(vram_mb) => {
                println!("✓ {:.1} MB", vram_mb);
                batch_tests.push(BatchTest {
                    batch_size: batch_size as u32,
                    vram_used_mb: vram_mb,
                    success: true,
                    error_message: None,
                });
            },
            Err(e) => {
                println!("✗ Failed: {}", e);
                batch_tests.push(BatchTest {
                    batch_size: batch_size as u32,
                    vram_used_mb: 0.0,
                    success: false,
                    error_message: Some(e.to_string()),
                });
                break; // Stop on first failure
            },
        }
    }

    let peak_vram = batch_tests
        .iter()
        .filter(|t| t.success)
        .map(|t| t.vram_used_mb)
        .max_by(|a, b| a.partial_cmp(b).unwrap())
        .unwrap_or(base_vram_mb);

    let max_safe_batch = batch_tests
        .iter()
        .filter(|t| t.success && (t.vram_used_mb / gpu_total_mb) < 0.8)
        .map(|t| t.batch_size)
        .max()
        .unwrap_or(1);

    let status = ProfileStatus::from_usage(peak_vram, gpu_total_mb);

    Ok(GpuModelProfile {
        model_name: "DQN".to_string(),
        base_vram_mb,
        peak_vram_mb: peak_vram,
        batch_size_tests: batch_tests,
        max_safe_batch_size: max_safe_batch,
        training_batch_size: std::cmp::min(max_safe_batch, 64),
        inference_batch_size: std::cmp::min(max_safe_batch, 128),
        parameter_count: (hidden_dim * input_dim)
            + hidden_dim
            + (hidden_dim * hidden_dim)
            + hidden_dim
            + (action_dim * hidden_dim)
            + action_dim,
        status,
    })
}

/// Profile PPO model VRAM usage (actor + critic)
fn profile_ppo_vram(device: &Device, gpu_total_mb: f64) -> Result<GpuModelProfile> {
    println!("📊 Profiling PPO VRAM usage...");

    let input_dim = 64;
    let hidden_dim = 256;
    let action_dim = 5;

    let baseline = query_gpu_memory()?;
    let base_vram_mb = baseline.used_mb;

    // Actor network
    let actor_l1 = Tensor::randn(0f32, 1.0, (hidden_dim, input_dim), device)?;
    let actor_l2 = Tensor::randn(0f32, 1.0, (hidden_dim, hidden_dim), device)?;
    let actor_out = Tensor::randn(0f32, 1.0, (action_dim, hidden_dim), device)?;

    // Critic network
    let critic_l1 = Tensor::randn(0f32, 1.0, (hidden_dim, input_dim), device)?;
    let critic_l2 = Tensor::randn(0f32, 1.0, (hidden_dim, hidden_dim), device)?;
    let critic_out = Tensor::randn(0f32, 1.0, (1, hidden_dim), device)?;

    std::thread::sleep(std::time::Duration::from_millis(100));
    let after_model = query_gpu_memory()?;
    println!(
        "  Model loaded: {:.1} MB VRAM",
        after_model.used_mb - base_vram_mb
    );

    let mut batch_tests = Vec::new();
    let batch_sizes = vec![1, 8, 16, 32, 64, 128, 256];

    for batch_size in batch_sizes {
        print!("  Testing batch size {}: ", batch_size);

        let result = (|| -> Result<f64> {
            let input = Tensor::randn(0f32, 1.0, (batch_size, input_dim), device)?;

            // Actor forward
            let _actor_out = input
                .matmul(&actor_l1.t()?)?
                .relu()?
                .matmul(&actor_l2.t()?)?
                .relu()?
                .matmul(&actor_out.t()?)?;

            // Critic forward
            let _critic_out = input
                .matmul(&critic_l1.t()?)?
                .relu()?
                .matmul(&critic_l2.t()?)?
                .relu()?
                .matmul(&critic_out.t()?)?;

            std::thread::sleep(std::time::Duration::from_millis(100));
            let mem = query_gpu_memory()?;
            Ok(mem.used_mb)
        })();

        match result {
            Ok(vram_mb) => {
                println!("✓ {:.1} MB", vram_mb);
                batch_tests.push(BatchTest {
                    batch_size: batch_size as u32,
                    vram_used_mb: vram_mb,
                    success: true,
                    error_message: None,
                });
            },
            Err(e) => {
                println!("✗ Failed: {}", e);
                batch_tests.push(BatchTest {
                    batch_size: batch_size as u32,
                    vram_used_mb: 0.0,
                    success: false,
                    error_message: Some(e.to_string()),
                });
                break;
            },
        }
    }

    let peak_vram = batch_tests
        .iter()
        .filter(|t| t.success)
        .map(|t| t.vram_used_mb)
        .max_by(|a, b| a.partial_cmp(b).unwrap())
        .unwrap_or(base_vram_mb);

    let max_safe_batch = batch_tests
        .iter()
        .filter(|t| t.success && (t.vram_used_mb / gpu_total_mb) < 0.8)
        .map(|t| t.batch_size)
        .max()
        .unwrap_or(1);

    let status = ProfileStatus::from_usage(peak_vram, gpu_total_mb);

    let actor_params =
        (hidden_dim * input_dim) + (hidden_dim * hidden_dim) + (action_dim * hidden_dim);
    let critic_params = (hidden_dim * input_dim) + (hidden_dim * hidden_dim) + (1 * hidden_dim);

    Ok(GpuModelProfile {
        model_name: "PPO".to_string(),
        base_vram_mb,
        peak_vram_mb: peak_vram,
        batch_size_tests: batch_tests,
        max_safe_batch_size: max_safe_batch,
        training_batch_size: std::cmp::min(max_safe_batch, 64),
        inference_batch_size: std::cmp::min(max_safe_batch, 128),
        parameter_count: actor_params + critic_params,
        status,
    })
}

/// Profile MAMBA-2 model VRAM usage
fn profile_mamba2_vram(device: &Device, gpu_total_mb: f64) -> Result<GpuModelProfile> {
    println!("📊 Profiling MAMBA-2 VRAM usage...");

    let d_model = 256;
    let n_layers = 4;

    let baseline = query_gpu_memory()?;
    let base_vram_mb = baseline.used_mb;

    // Simplified state space layers
    let mut layers = Vec::new();
    for _ in 0..n_layers {
        let a = Tensor::randn(0f32, 1.0, (d_model, d_model), device)?;
        let b = Tensor::randn(0f32, 1.0, (d_model, d_model), device)?;
        let c = Tensor::randn(0f32, 1.0, (d_model, d_model), device)?;
        layers.push((a, b, c));
    }

    std::thread::sleep(std::time::Duration::from_millis(100));
    let after_model = query_gpu_memory()?;
    println!(
        "  Model loaded: {:.1} MB VRAM",
        after_model.used_mb - base_vram_mb
    );

    let mut batch_tests = Vec::new();
    let batch_sizes = vec![1, 4, 8, 16, 32, 64];
    let seq_len = 64;

    for batch_size in batch_sizes {
        print!("  Testing batch size {}: ", batch_size);

        let result = (|| -> Result<f64> {
            let input = Tensor::randn(0f32, 1.0, (batch_size, seq_len, d_model), device)?;
            let _state = Tensor::randn(0f32, 1.0, (batch_size, d_model, d_model), device)?;

            std::thread::sleep(std::time::Duration::from_millis(100));
            let mem = query_gpu_memory()?;
            Ok(mem.used_mb)
        })();

        match result {
            Ok(vram_mb) => {
                println!("✓ {:.1} MB", vram_mb);
                batch_tests.push(BatchTest {
                    batch_size: batch_size as u32,
                    vram_used_mb: vram_mb,
                    success: true,
                    error_message: None,
                });
            },
            Err(e) => {
                println!("✗ Failed: {}", e);
                batch_tests.push(BatchTest {
                    batch_size: batch_size as u32,
                    vram_used_mb: 0.0,
                    success: false,
                    error_message: Some(e.to_string()),
                });
                break;
            },
        }
    }

    let peak_vram = batch_tests
        .iter()
        .filter(|t| t.success)
        .map(|t| t.vram_used_mb)
        .max_by(|a, b| a.partial_cmp(b).unwrap())
        .unwrap_or(base_vram_mb);

    let max_safe_batch = batch_tests
        .iter()
        .filter(|t| t.success && (t.vram_used_mb / gpu_total_mb) < 0.8)
        .map(|t| t.batch_size)
        .max()
        .unwrap_or(1);

    let status = ProfileStatus::from_usage(peak_vram, gpu_total_mb);

    Ok(GpuModelProfile {
        model_name: "MAMBA-2".to_string(),
        base_vram_mb,
        peak_vram_mb: peak_vram,
        batch_size_tests: batch_tests,
        max_safe_batch_size: max_safe_batch,
        training_batch_size: std::cmp::min(max_safe_batch, 32),
        inference_batch_size: std::cmp::min(max_safe_batch, 64),
        parameter_count: n_layers * (d_model * d_model * 3),
        status,
    })
}

/// Profile TFT model VRAM usage (memory-intensive)
fn profile_tft_vram(device: &Device, gpu_total_mb: f64) -> Result<GpuModelProfile> {
    println!("📊 Profiling TFT VRAM usage...");

    let d_model = 512;
    let n_heads = 8;
    let n_layers = 6;

    let baseline = query_gpu_memory()?;
    let base_vram_mb = baseline.used_mb;

    // Attention layers (most memory-intensive)
    let mut layers = Vec::new();
    for _ in 0..n_layers {
        let q_proj = Tensor::randn(0f32, 1.0, (d_model, d_model), device)?;
        let k_proj = Tensor::randn(0f32, 1.0, (d_model, d_model), device)?;
        let v_proj = Tensor::randn(0f32, 1.0, (d_model, d_model), device)?;
        let o_proj = Tensor::randn(0f32, 1.0, (d_model, d_model), device)?;
        layers.push((q_proj, k_proj, v_proj, o_proj));
    }

    std::thread::sleep(std::time::Duration::from_millis(100));
    let after_model = query_gpu_memory()?;
    println!(
        "  Model loaded: {:.1} MB VRAM",
        after_model.used_mb - base_vram_mb
    );

    let mut batch_tests = Vec::new();
    let batch_sizes = vec![1, 2, 4, 8, 16, 32];
    let seq_len = 64;

    for batch_size in batch_sizes {
        print!("  Testing batch size {}: ", batch_size);

        let result = (|| -> Result<f64> {
            let input = Tensor::randn(0f32, 1.0, (batch_size, seq_len, d_model), device)?;

            // Attention is the memory bottleneck (seq_len x seq_len attention matrix)
            let q = Tensor::randn(
                0f32,
                1.0,
                (batch_size, n_heads, seq_len, d_model / n_heads),
                device,
            )?;
            let k = Tensor::randn(
                0f32,
                1.0,
                (batch_size, n_heads, seq_len, d_model / n_heads),
                device,
            )?;
            let _attn_scores =
                Tensor::randn(0f32, 1.0, (batch_size, n_heads, seq_len, seq_len), device)?;

            std::thread::sleep(std::time::Duration::from_millis(100));
            let mem = query_gpu_memory()?;
            Ok(mem.used_mb)
        })();

        match result {
            Ok(vram_mb) => {
                println!("✓ {:.1} MB", vram_mb);
                batch_tests.push(BatchTest {
                    batch_size: batch_size as u32,
                    vram_used_mb: vram_mb,
                    success: true,
                    error_message: None,
                });
            },
            Err(e) => {
                println!("✗ Failed: {}", e);
                batch_tests.push(BatchTest {
                    batch_size: batch_size as u32,
                    vram_used_mb: 0.0,
                    success: false,
                    error_message: Some(e.to_string()),
                });
                break;
            },
        }
    }

    let peak_vram = batch_tests
        .iter()
        .filter(|t| t.success)
        .map(|t| t.vram_used_mb)
        .max_by(|a, b| a.partial_cmp(b).unwrap())
        .unwrap_or(base_vram_mb);

    let max_safe_batch = batch_tests
        .iter()
        .filter(|t| t.success && (t.vram_used_mb / gpu_total_mb) < 0.8)
        .map(|t| t.batch_size)
        .max()
        .unwrap_or(1);

    let status = ProfileStatus::from_usage(peak_vram, gpu_total_mb);

    Ok(GpuModelProfile {
        model_name: "TFT".to_string(),
        base_vram_mb,
        peak_vram_mb: peak_vram,
        batch_size_tests: batch_tests,
        max_safe_batch_size: max_safe_batch,
        training_batch_size: std::cmp::min(max_safe_batch, 8),
        inference_batch_size: std::cmp::min(max_safe_batch, 16),
        parameter_count: n_layers * (d_model * d_model * 4),
        status,
    })
}

/// Profile Liquid NN VRAM usage
fn profile_liquid_vram(device: &Device, gpu_total_mb: f64) -> Result<GpuModelProfile> {
    println!("📊 Profiling Liquid NN VRAM usage...");

    let input_dim = 64;
    let hidden_dim = 256;
    let output_dim = 5;

    let baseline = query_gpu_memory()?;
    let base_vram_mb = baseline.used_mb;

    // Liquid NN weights
    let w_in = Tensor::randn(0f32, 1.0, (hidden_dim, input_dim), device)?;
    let w_rec = Tensor::randn(0f32, 1.0, (hidden_dim, hidden_dim), device)?;
    let w_out = Tensor::randn(0f32, 1.0, (output_dim, hidden_dim), device)?;
    let tau = Tensor::randn(0f32, 1.0, (hidden_dim,), device)?;

    std::thread::sleep(std::time::Duration::from_millis(100));
    let after_model = query_gpu_memory()?;
    println!(
        "  Model loaded: {:.1} MB VRAM",
        after_model.used_mb - base_vram_mb
    );

    let mut batch_tests = Vec::new();
    let batch_sizes = vec![1, 8, 16, 32, 64, 128, 256];

    for batch_size in batch_sizes {
        print!("  Testing batch size {}: ", batch_size);

        let result = (|| -> Result<f64> {
            let input = Tensor::randn(0f32, 1.0, (batch_size, input_dim), device)?;
            let hidden = Tensor::randn(0f32, 1.0, (batch_size, hidden_dim), device)?;

            // ODE solver requires multiple intermediate states
            let _intermediates = Tensor::randn(0f32, 1.0, (batch_size, hidden_dim, 10), device)?;

            std::thread::sleep(std::time::Duration::from_millis(100));
            let mem = query_gpu_memory()?;
            Ok(mem.used_mb)
        })();

        match result {
            Ok(vram_mb) => {
                println!("✓ {:.1} MB", vram_mb);
                batch_tests.push(BatchTest {
                    batch_size: batch_size as u32,
                    vram_used_mb: vram_mb,
                    success: true,
                    error_message: None,
                });
            },
            Err(e) => {
                println!("✗ Failed: {}", e);
                batch_tests.push(BatchTest {
                    batch_size: batch_size as u32,
                    vram_used_mb: 0.0,
                    success: false,
                    error_message: Some(e.to_string()),
                });
                break;
            },
        }
    }

    let peak_vram = batch_tests
        .iter()
        .filter(|t| t.success)
        .map(|t| t.vram_used_mb)
        .max_by(|a, b| a.partial_cmp(b).unwrap())
        .unwrap_or(base_vram_mb);

    let max_safe_batch = batch_tests
        .iter()
        .filter(|t| t.success && (t.vram_used_mb / gpu_total_mb) < 0.8)
        .map(|t| t.batch_size)
        .max()
        .unwrap_or(1);

    let status = ProfileStatus::from_usage(peak_vram, gpu_total_mb);

    Ok(GpuModelProfile {
        model_name: "Liquid NN".to_string(),
        base_vram_mb,
        peak_vram_mb: peak_vram,
        batch_size_tests: batch_tests,
        max_safe_batch_size: max_safe_batch,
        training_batch_size: std::cmp::min(max_safe_batch, 64),
        inference_batch_size: std::cmp::min(max_safe_batch, 128),
        parameter_count: (hidden_dim * input_dim)
            + (hidden_dim * hidden_dim)
            + (output_dim * hidden_dim)
            + hidden_dim,
        status,
    })
}

/// Generate comprehensive markdown report
fn generate_report(
    profiles: &[GpuModelProfile],
    gpu_snapshot: &GpuMemorySnapshot,
    output_path: &PathBuf,
) -> Result<()> {
    let mut file = File::create(output_path)?;

    writeln!(file, "# GPU Memory Profile Report - RTX 3050 Ti (4GB VRAM)")?;
    writeln!(file)?;
    writeln!(
        file,
        "**Generated**: {}",
        Utc::now().format("%Y-%m-%d %H:%M:%S UTC")
    )?;
    writeln!(file, "**GPU**: NVIDIA GeForce RTX 3050 Ti Laptop")?;
    writeln!(
        file,
        "**VRAM**: {:.0} MB total, {:.0} MB free at start",
        gpu_snapshot.total_mb, gpu_snapshot.free_mb
    )?;
    writeln!(file)?;
    writeln!(file, "---")?;
    writeln!(file)?;

    // Executive Summary
    writeln!(file, "## Executive Summary")?;
    writeln!(file)?;
    writeln!(file, "This report profiles GPU VRAM usage for all ML models using direct `nvidia-smi` measurements.")?;
    writeln!(file)?;

    for profile in profiles {
        writeln!(file, "- **{}**: {:.1} MB peak VRAM, batch size {} (training), batch size {} (inference) - {} {}",
                 profile.model_name, profile.peak_vram_mb,
                 profile.training_batch_size, profile.inference_batch_size,
                 profile.status.emoji(),
                 match profile.status {
                     ProfileStatus::Safe => "Safe",
                     ProfileStatus::Tight => "Tight",
                     ProfileStatus::Critical => "Critical",
                     ProfileStatus::Failed => "Failed",
                 })?;
    }
    writeln!(file)?;

    // Detailed Profiles
    writeln!(file, "---")?;
    writeln!(file)?;
    writeln!(file, "## Detailed Model Profiles")?;
    writeln!(file)?;

    for profile in profiles {
        writeln!(file, "### {}", profile.model_name)?;
        writeln!(file)?;
        writeln!(file, "- **Parameters**: {}", profile.parameter_count)?;
        writeln!(file, "- **Base VRAM**: {:.1} MB", profile.base_vram_mb)?;
        writeln!(file, "- **Peak VRAM**: {:.1} MB", profile.peak_vram_mb)?;
        writeln!(
            file,
            "- **Status**: {} {:?}",
            profile.status.emoji(),
            profile.status
        )?;
        writeln!(
            file,
            "- **Max Safe Batch Size**: {}",
            profile.max_safe_batch_size
        )?;
        writeln!(
            file,
            "- **Training Batch Size**: {}",
            profile.training_batch_size
        )?;
        writeln!(
            file,
            "- **Inference Batch Size**: {}",
            profile.inference_batch_size
        )?;
        writeln!(file)?;

        writeln!(file, "#### Batch Size Tests")?;
        writeln!(file)?;
        writeln!(file, "| Batch Size | VRAM (MB) | Status |")?;
        writeln!(file, "|------------|-----------|--------|")?;

        for test in &profile.batch_size_tests {
            if test.success {
                let usage_pct = (test.vram_used_mb / gpu_snapshot.total_mb) * 100.0;
                writeln!(
                    file,
                    "| {} | {:.1} ({:.0}%) | ✅ Success |",
                    test.batch_size, test.vram_used_mb, usage_pct
                )?;
            } else {
                writeln!(file, "| {} | OOM | ❌ Failed |", test.batch_size)?;
            }
        }
        writeln!(file)?;
    }

    // Memory Budget
    writeln!(file, "---")?;
    writeln!(file)?;
    writeln!(file, "## Memory Budget Allocation")?;
    writeln!(file)?;
    writeln!(file, "### Training (Single Model)")?;
    writeln!(file)?;
    writeln!(file, "| Model | Peak VRAM | Training Batch | Status |")?;
    writeln!(file, "|-------|-----------|----------------|--------|")?;

    for profile in profiles {
        writeln!(
            file,
            "| {} | {:.1} MB | {} | {} |",
            profile.model_name,
            profile.peak_vram_mb,
            profile.training_batch_size,
            profile.status.emoji()
        )?;
    }
    writeln!(file)?;

    writeln!(file, "### Inference (Multi-Model Ensemble)")?;
    writeln!(file)?;

    let total_vram: f64 = profiles.iter().map(|p| p.base_vram_mb).sum();
    let can_load_all = total_vram < gpu_snapshot.total_mb * 0.9;

    writeln!(
        file,
        "- **Total VRAM for all models**: {:.1} MB",
        total_vram
    )?;
    writeln!(
        file,
        "- **Available VRAM**: {:.1} MB",
        gpu_snapshot.total_mb
    )?;
    writeln!(
        file,
        "- **Can load all models**: {}",
        if can_load_all {
            "✅ Yes"
        } else {
            "❌ No (use hot-swapping)"
        }
    )?;
    writeln!(file)?;

    if !can_load_all {
        let models_can_fit =
            (gpu_snapshot.total_mb * 0.9 / (total_vram / profiles.len() as f64)).floor() as u32;
        writeln!(
            file,
            "- **Simultaneous models**: Up to {} models recommended",
            models_can_fit
        )?;
        writeln!(
            file,
            "- **Recommendation**: Use LRU cache with hot-swapping"
        )?;
    }
    writeln!(file)?;

    // Recommendations
    writeln!(file, "---")?;
    writeln!(file)?;
    writeln!(file, "## Recommendations")?;
    writeln!(file)?;
    writeln!(file, "### Training")?;
    writeln!(file)?;
    writeln!(
        file,
        "1. **Train one model at a time** - Use recommended batch sizes above"
    )?;
    writeln!(
        file,
        "2. **Monitor VRAM** - Run `watch -n1 nvidia-smi` during training"
    )?;
    writeln!(
        file,
        "3. **Use gradient accumulation** for TFT model (small batch size)"
    )?;
    writeln!(
        file,
        "4. **Enable mixed precision (FP16)** to reduce VRAM by ~40%"
    )?;
    writeln!(
        file,
        "5. **Clear CUDA cache** between model switches: `torch.cuda.empty_cache()`"
    )?;
    writeln!(file)?;
    writeln!(file, "### Inference")?;
    writeln!(file)?;

    if can_load_all {
        writeln!(
            file,
            "1. **All models can be loaded simultaneously** for ensemble inference"
        )?;
        writeln!(
            file,
            "2. **Use batch inference** with recommended batch sizes"
        )?;
    } else {
        writeln!(file, "1. **Implement model hot-swapping** with LRU cache")?;
        writeln!(file, "2. **Load models on-demand** for predictions")?;
        writeln!(file, "3. **Consider FP16 quantization** to fit more models")?;
    }
    writeln!(file)?;

    // Expected vs Actual
    writeln!(file, "---")?;
    writeln!(file)?;
    writeln!(file, "## Expected vs Actual VRAM Usage")?;
    writeln!(file)?;
    writeln!(
        file,
        "| Model | Expected Range (MB) | Actual (MB) | Status |"
    )?;
    writeln!(
        file,
        "|-------|---------------------|-------------|--------|"
    )?;

    let comparisons = vec![
        ("DQN", DQN_RANGE_MB),
        ("PPO", PPO_RANGE_MB),
        ("MAMBA-2", MAMBA2_RANGE_MB),
        ("TFT", TFT_RANGE_MB),
        ("Liquid NN", LIQUID_RANGE_MB),
    ];

    for (name, (min_exp, max_exp)) in &comparisons {
        if let Some(profile) = profiles.iter().find(|p| p.model_name == *name) {
            let actual = profile.peak_vram_mb;
            let status_str = if actual >= *min_exp && actual <= *max_exp {
                "✅ Within range"
            } else if actual < *min_exp {
                "⚠️ Lower"
            } else {
                "⚠️ Higher"
            };

            writeln!(
                file,
                "| {} | {:.0}-{:.0} | {:.1} | {} |",
                name, min_exp, max_exp, actual, status_str
            )?;
        }
    }
    writeln!(file)?;

    writeln!(file, "---")?;
    writeln!(file)?;
    writeln!(file, "**Agent**: 133 (GPU Memory Profiling)")?;
    writeln!(
        file,
        "**Command**: `cargo run -p ml --example gpu_memory_benchmark --release --features cuda`"
    )?;

    Ok(())
}

fn main() -> Result<()> {
    println!("🚀 GPU Memory Profiling for RTX 3050 Ti (4GB VRAM)\n");

    // Verify CUDA availability
    let device = match Device::new_cuda(0) {
        Ok(dev) => dev,
        Err(_) => {
            eprintln!("❌ CUDA device not available");
            eprintln!("   Ensure CUDA is installed and GPU is accessible");
            std::process::exit(1);
        },
    };

    println!("✓ CUDA device initialized: {:?}\n", device);

    // Get GPU memory baseline
    let gpu_snapshot = query_gpu_memory()?;
    println!("📊 GPU Memory Baseline:");
    println!("   Total: {:.1} MB", gpu_snapshot.total_mb);
    println!("   Free: {:.1} MB", gpu_snapshot.free_mb);
    println!("   Used: {:.1} MB", gpu_snapshot.used_mb);
    println!("   Utilization: {:.0}%\n", gpu_snapshot.utilization_percent);

    // Profile each model
    let mut profiles = Vec::new();

    // DQN
    match profile_dqn_vram(&device, gpu_snapshot.total_mb) {
        Ok(profile) => {
            println!(
                "✅ DQN: {:.1} MB peak, batch size {}\n",
                profile.peak_vram_mb, profile.max_safe_batch_size
            );
            profiles.push(profile);
        },
        Err(e) => eprintln!("❌ DQN profiling failed: {}\n", e),
    }

    // PPO
    match profile_ppo_vram(&device, gpu_snapshot.total_mb) {
        Ok(profile) => {
            println!(
                "✅ PPO: {:.1} MB peak, batch size {}\n",
                profile.peak_vram_mb, profile.max_safe_batch_size
            );
            profiles.push(profile);
        },
        Err(e) => eprintln!("❌ PPO profiling failed: {}\n", e),
    }

    // MAMBA-2
    match profile_mamba2_vram(&device, gpu_snapshot.total_mb) {
        Ok(profile) => {
            println!(
                "✅ MAMBA-2: {:.1} MB peak, batch size {}\n",
                profile.peak_vram_mb, profile.max_safe_batch_size
            );
            profiles.push(profile);
        },
        Err(e) => eprintln!("❌ MAMBA-2 profiling failed: {}\n", e),
    }

    // TFT
    match profile_tft_vram(&device, gpu_snapshot.total_mb) {
        Ok(profile) => {
            println!(
                "✅ TFT: {:.1} MB peak, batch size {}\n",
                profile.peak_vram_mb, profile.max_safe_batch_size
            );
            profiles.push(profile);
        },
        Err(e) => eprintln!("❌ TFT profiling failed: {}\n", e),
    }

    // Liquid NN
    match profile_liquid_vram(&device, gpu_snapshot.total_mb) {
        Ok(profile) => {
            println!(
                "✅ Liquid NN: {:.1} MB peak, batch size {}\n",
                profile.peak_vram_mb, profile.max_safe_batch_size
            );
            profiles.push(profile);
        },
        Err(e) => eprintln!("❌ Liquid NN profiling failed: {}\n", e),
    }

    // Generate report
    let output_path = PathBuf::from("/home/jgrusewski/Work/foxhunt/GPU_MEMORY_PROFILE_REPORT.md");
    generate_report(&profiles, &gpu_snapshot, &output_path)?;

    println!("\n✅ GPU memory profiling complete!");
    println!("📄 Report saved to: {}", output_path.display());
    println!("\n🎯 Summary:");

    for profile in &profiles {
        println!(
            "   {} - {:.1} MB peak VRAM (batch size: {})",
            profile.model_name, profile.peak_vram_mb, profile.max_safe_batch_size
        );
    }

    let total_vram: f64 = profiles.iter().map(|p| p.base_vram_mb).sum();
    println!(
        "\n   Total VRAM for all models: {:.1} MB / {:.1} MB available",
        total_vram, gpu_snapshot.total_mb
    );

    if total_vram > gpu_snapshot.total_mb * 0.9 {
        println!("   ⚠️  Use model hot-swapping for ensemble inference");
    } else {
        println!("   ✅ All models can be loaded simultaneously");
    }

    Ok(())
}