//! GPU Batch Size Optimization for RTX 3050 Ti (4GB VRAM)
//!
//! Tests multiple batch sizes for TFT, MAMBA-2, and Liquid models to find
//! optimal configurations that maximize throughput while staying under 4GB VRAM.
//!
//! ## Testing Strategy
//!
//! 1. **TFT**: Test batch sizes [16, 32, 64, 128]
//! 2. **MAMBA-2**: Test batch sizes [8, 16, 32]
//! 3. **Liquid**: Test batch sizes [16, 32, 64]
//! 4. Monitor VRAM usage with nvidia-smi integration
//! 5. Find optimal batch size (max throughput, <4GB VRAM)
//!
//! ## Usage
//!
//! ```bash
//! cargo run -p ml --example optimize_batch_sizes --release
//! ```
//!
//! ## Output
//!
//! Generates `BATCH_SIZE_OPTIMIZATION_REPORT.md` with:
//! - VRAM usage per model/batch size
//! - Throughput measurements
//! - Recommended optimal batch sizes
//! - Updated configuration snippets

use candle_core::{DType, Device, Tensor};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::fs::File;
use std::io::Write;
use std::process::Command;
use std::time::Instant;
use sysinfo::{ProcessExt, System, SystemExt};

// Import model types
use ml::liquid::network::{LiquidNetwork, LiquidNetworkConfig};
use ml::mamba::{Mamba2Config, Mamba2SSM};
use ml::tft::{TFTConfig, TemporalFusionTransformer};

const VRAM_LIMIT_GB: f32 = 4.0;
const WARMUP_ITERATIONS: usize = 5;
const BENCHMARK_ITERATIONS: usize = 20;

#[derive(Debug, Clone, Serialize, Deserialize)]
struct BatchSizeResult {
    model_name: String,
    batch_size: usize,
    vram_used_mb: f32,
    vram_percent: f32,
    throughput_samples_per_sec: f32,
    latency_ms: f32,
    oom_occurred: bool,
    recommended: bool,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
struct OptimizationReport {
    timestamp: String,
    gpu_name: String,
    vram_total_gb: f32,
    results: Vec<BatchSizeResult>,
    recommendations: HashMap<String, usize>,
    notes: Vec<String>,
}

/// Query NVIDIA GPU VRAM usage in MB
fn query_nvidia_vram() -> Result<f32, Box<dyn std::error::Error>> {
    let output = Command::new("nvidia-smi")
        .args(&["--query-gpu=memory.used", "--format=csv,noheader,nounits"])
        .output()?;

    let vram_str = String::from_utf8(output.stdout)?;
    let vram_mb: f32 = vram_str.trim().parse()?;
    Ok(vram_mb)
}

/// Query NVIDIA GPU name
fn query_gpu_name() -> Result<String, Box<dyn std::error::Error>> {
    let output = Command::new("nvidia-smi")
        .args(&["--query-gpu=name", "--format=csv,noheader"])
        .output()?;

    Ok(String::from_utf8(output.stdout)?.trim().to_string())
}

/// Query total VRAM in GB
fn query_total_vram_gb() -> Result<f32, Box<dyn std::error::Error>> {
    let output = Command::new("nvidia-smi")
        .args(&["--query-gpu=memory.total", "--format=csv,noheader,nounits"])
        .output()?;

    let vram_mb: f32 = String::from_utf8(output.stdout)?.trim().parse()?;
    Ok(vram_mb / 1024.0) // Convert to GB
}

/// Benchmark TFT model with specific batch size
fn benchmark_tft_batch(
    batch_size: usize,
    device: &Device,
) -> Result<BatchSizeResult, Box<dyn std::error::Error>> {
    println!("  Testing TFT with batch_size={}", batch_size);

    // Create TFT configuration
    let config = TFTConfig {
        input_dim: 64,
        hidden_dim: 256,
        num_heads: 8,
        num_layers: 4,
        prediction_horizon: 10,
        sequence_length: 50,
        num_quantiles: 9,
        num_static_features: 5,
        num_known_features: 10,
        num_unknown_features: 20,
        learning_rate: 1e-3,
        batch_size,
        dropout_rate: 0.1,
        l2_regularization: 1e-4,
        use_flash_attention: true,
        mixed_precision: false,
        memory_efficient: true,
        max_inference_latency_us: 50,
        target_throughput_pps: 100_000,
    };

    let mut model = match TemporalFusionTransformer::new(config.clone()) {
        Ok(m) => m,
        Err(e) => {
            return Ok(BatchSizeResult {
                model_name: "TFT".to_string(),
                batch_size,
                vram_used_mb: 0.0,
                vram_percent: 0.0,
                throughput_samples_per_sec: 0.0,
                latency_ms: 0.0,
                oom_occurred: true,
                recommended: false,
            });
        },
    };

    // Prepare batch inputs
    let static_features = vec![0.0f32; config.num_static_features * batch_size];
    let historical_features =
        vec![0.0f32; config.sequence_length * config.num_unknown_features * batch_size];
    let future_features =
        vec![0.0f32; config.prediction_horizon * config.num_known_features * batch_size];

    // Warmup
    for _ in 0..WARMUP_ITERATIONS {
        let _ = model.predict_fast(&static_features, &historical_features, &future_features);
    }

    // Clear GPU cache
    std::thread::sleep(std::time::Duration::from_millis(500));

    // Measure baseline VRAM
    let baseline_vram = query_nvidia_vram()?;

    // Benchmark iterations
    let start = Instant::now();
    for _ in 0..BENCHMARK_ITERATIONS {
        match model.predict_fast(&static_features, &historical_features, &future_features) {
            Ok(_) => {},
            Err(_) => {
                return Ok(BatchSizeResult {
                    model_name: "TFT".to_string(),
                    batch_size,
                    vram_used_mb: 0.0,
                    vram_percent: 0.0,
                    throughput_samples_per_sec: 0.0,
                    latency_ms: 0.0,
                    oom_occurred: true,
                    recommended: false,
                });
            },
        }
    }
    let elapsed = start.elapsed();

    // Measure peak VRAM
    let peak_vram = query_nvidia_vram()?;
    let vram_used_mb = peak_vram - baseline_vram;
    let total_vram_gb = query_total_vram_gb()?;
    let vram_percent = (peak_vram / (total_vram_gb * 1024.0)) * 100.0;

    let latency_ms = elapsed.as_secs_f32() * 1000.0 / BENCHMARK_ITERATIONS as f32;
    let throughput = (batch_size * BENCHMARK_ITERATIONS) as f32 / elapsed.as_secs_f32();

    let recommended = vram_percent < 90.0 && !false;

    Ok(BatchSizeResult {
        model_name: "TFT".to_string(),
        batch_size,
        vram_used_mb,
        vram_percent,
        throughput_samples_per_sec: throughput,
        latency_ms,
        oom_occurred: false,
        recommended,
    })
}

/// Benchmark MAMBA-2 model with specific batch size
fn benchmark_mamba_batch(
    batch_size: usize,
    device: &Device,
) -> Result<BatchSizeResult, Box<dyn std::error::Error>> {
    println!("  Testing MAMBA-2 with batch_size={}", batch_size);

    let config = Mamba2Config {
        d_model: 256,
        d_state: 64,
        d_head: 32,
        num_heads: 8,
        expand: 2,
        num_layers: 4,
        dropout: 0.1,
        use_ssd: true,
        use_selective_state: true,
        hardware_aware: true,
        target_latency_us: 5,
        max_seq_len: 512,
        learning_rate: 1e-3,
        weight_decay: 1e-4,
        grad_clip: 1.0,
        warmup_steps: 1000,
        batch_size,
        seq_len: 256,
    };

    let mut model = match Mamba2SSM::new(config.clone(), device) {
        Ok(m) => m,
        Err(e) => {
            return Ok(BatchSizeResult {
                model_name: "MAMBA-2".to_string(),
                batch_size,
                vram_used_mb: 0.0,
                vram_percent: 0.0,
                throughput_samples_per_sec: 0.0,
                latency_ms: 0.0,
                oom_occurred: true,
                recommended: false,
            });
        },
    };

    // Prepare batch input
    let input_tensor = match Tensor::randn(0f32, 1f32, (batch_size, config.d_model), device) {
        Ok(t) => t,
        Err(_) => {
            return Ok(BatchSizeResult {
                model_name: "MAMBA-2".to_string(),
                batch_size,
                vram_used_mb: 0.0,
                vram_percent: 0.0,
                throughput_samples_per_sec: 0.0,
                latency_ms: 0.0,
                oom_occurred: true,
                recommended: false,
            });
        },
    };

    // Warmup
    for _ in 0..WARMUP_ITERATIONS {
        let _ = model.forward(&input_tensor);
    }

    std::thread::sleep(std::time::Duration::from_millis(500));

    let baseline_vram = query_nvidia_vram()?;

    // Benchmark
    let start = Instant::now();
    for _ in 0..BENCHMARK_ITERATIONS {
        match model.forward(&input_tensor) {
            Ok(_) => {},
            Err(_) => {
                return Ok(BatchSizeResult {
                    model_name: "MAMBA-2".to_string(),
                    batch_size,
                    vram_used_mb: 0.0,
                    vram_percent: 0.0,
                    throughput_samples_per_sec: 0.0,
                    latency_ms: 0.0,
                    oom_occurred: true,
                    recommended: false,
                });
            },
        }
    }
    let elapsed = start.elapsed();

    let peak_vram = query_nvidia_vram()?;
    let vram_used_mb = peak_vram - baseline_vram;
    let total_vram_gb = query_total_vram_gb()?;
    let vram_percent = (peak_vram / (total_vram_gb * 1024.0)) * 100.0;

    let latency_ms = elapsed.as_secs_f32() * 1000.0 / BENCHMARK_ITERATIONS as f32;
    let throughput = (batch_size * BENCHMARK_ITERATIONS) as f32 / elapsed.as_secs_f32();

    let recommended = vram_percent < 90.0 && !false;

    Ok(BatchSizeResult {
        model_name: "MAMBA-2".to_string(),
        batch_size,
        vram_used_mb,
        vram_percent,
        throughput_samples_per_sec: throughput,
        latency_ms,
        oom_occurred: false,
        recommended,
    })
}

/// Benchmark Liquid model with specific batch size
fn benchmark_liquid_batch(
    batch_size: usize,
) -> Result<BatchSizeResult, Box<dyn std::error::Error>> {
    println!("  Testing Liquid with batch_size={}", batch_size);

    let config = LiquidNetworkConfig {
        input_dim: 256,
        hidden_dim: 512,
        output_dim: 3,
        num_layers: 4,
        cell_type: ml::liquid::cells::CellType::LTC,
        activation: ml::liquid::activation::ActivationType::Tanh,
        solver: ml::liquid::ode_solvers::SolverType::Euler,
        time_step: 0.001,
        inference_steps: 10,
    };

    let model = match LiquidNetwork::new(config.clone()) {
        Ok(m) => m,
        Err(e) => {
            return Ok(BatchSizeResult {
                model_name: "Liquid".to_string(),
                batch_size,
                vram_used_mb: 0.0,
                vram_percent: 0.0,
                throughput_samples_per_sec: 0.0,
                latency_ms: 0.0,
                oom_occurred: true,
                recommended: false,
            });
        },
    };

    // Prepare batch input (Note: Liquid uses CPU, no GPU VRAM)
    let input = vec![ml::liquid::FixedPoint::zero(); config.input_dim * batch_size];

    // Warmup
    for _ in 0..WARMUP_ITERATIONS {
        let _ = model.forward(&input[..config.input_dim]);
    }

    std::thread::sleep(std::time::Duration::from_millis(500));

    let baseline_vram = query_nvidia_vram().unwrap_or(0.0);

    // Benchmark
    let start = Instant::now();
    for i in 0..BENCHMARK_ITERATIONS {
        let offset = (i % batch_size) * config.input_dim;
        match model.forward(&input[offset..offset + config.input_dim]) {
            Ok(_) => {},
            Err(_) => {
                return Ok(BatchSizeResult {
                    model_name: "Liquid".to_string(),
                    batch_size,
                    vram_used_mb: 0.0,
                    vram_percent: 0.0,
                    throughput_samples_per_sec: 0.0,
                    latency_ms: 0.0,
                    oom_occurred: true,
                    recommended: false,
                });
            },
        }
    }
    let elapsed = start.elapsed();

    let peak_vram = query_nvidia_vram().unwrap_or(0.0);
    let vram_used_mb = peak_vram - baseline_vram;
    let total_vram_gb = query_total_vram_gb().unwrap_or(4.0);
    let vram_percent = (peak_vram / (total_vram_gb * 1024.0)) * 100.0;

    let latency_ms = elapsed.as_secs_f32() * 1000.0 / BENCHMARK_ITERATIONS as f32;
    let throughput = (batch_size * BENCHMARK_ITERATIONS) as f32 / elapsed.as_secs_f32();

    // Liquid is CPU-based, always safe
    let recommended = true;

    Ok(BatchSizeResult {
        model_name: "Liquid".to_string(),
        batch_size,
        vram_used_mb,
        vram_percent,
        throughput_samples_per_sec: throughput,
        latency_ms,
        oom_occurred: false,
        recommended,
    })
}

/// Generate markdown report
fn generate_report(report: &OptimizationReport) -> String {
    let mut md = String::new();

    md.push_str("# GPU Batch Size Optimization Report\n\n");
    md.push_str(&format!("**Generated**: {}\n", report.timestamp));
    md.push_str(&format!("**GPU**: {}\n", report.gpu_name));
    md.push_str(&format!("**VRAM**: {:.1} GB\n\n", report.vram_total_gb));

    md.push_str("## Optimization Summary\n\n");
    md.push_str("| Model | Recommended Batch Size | VRAM Usage | Throughput |\n");
    md.push_str("|-------|------------------------|------------|------------|\n");
    for (model, batch) in &report.recommendations {
        let result = report
            .results
            .iter()
            .find(|r| r.model_name == *model && r.batch_size == *batch)
            .unwrap();
        md.push_str(&format!(
            "| {} | {} | {:.0} MB ({:.1}%) | {:.1} samples/sec |\n",
            model,
            batch,
            result.vram_used_mb,
            result.vram_percent,
            result.throughput_samples_per_sec
        ));
    }
    md.push_str("\n");

    md.push_str("## Detailed Results\n\n");

    // Group by model
    let mut by_model: HashMap<String, Vec<&BatchSizeResult>> = HashMap::new();
    for result in &report.results {
        by_model
            .entry(result.model_name.clone())
            .or_insert_with(Vec::new)
            .push(result);
    }

    for (model, results) in by_model {
        md.push_str(&format!("### {} Model\n\n", model));
        md.push_str("| Batch Size | VRAM (MB) | VRAM % | Latency (ms) | Throughput | Status |\n");
        md.push_str("|------------|-----------|--------|--------------|------------|--------|\n");

        for result in results {
            let status = if result.oom_occurred {
                "OOM"
            } else if result.recommended {
                "✅ Optimal"
            } else {
                "OK"
            };

            md.push_str(&format!(
                "| {} | {:.0} | {:.1}% | {:.2} | {:.1}/sec | {} |\n",
                result.batch_size,
                result.vram_used_mb,
                result.vram_percent,
                result.latency_ms,
                result.throughput_samples_per_sec,
                status
            ));
        }
        md.push_str("\n");
    }

    md.push_str("## Updated Configuration Snippets\n\n");

    for (model, batch) in &report.recommendations {
        md.push_str(&format!("### {} Configuration\n\n", model));
        md.push_str("```rust\n");
        match model.as_str() {
            "TFT" => {
                md.push_str(&format!(
                    "TFTConfig {{\n    batch_size: {},\n    // ... other fields\n}}\n",
                    batch
                ));
            },
            "MAMBA-2" => {
                md.push_str(&format!(
                    "Mamba2Config {{\n    batch_size: {},\n    // ... other fields\n}}\n",
                    batch
                ));
            },
            "Liquid" => {
                md.push_str(&format!(
                    "// Note: Liquid processes samples sequentially\n// Batch size {} tested for CPU efficiency\n",
                    batch
                ));
            },
            _ => {},
        }
        md.push_str("```\n\n");
    }

    md.push_str("## Notes\n\n");
    for note in &report.notes {
        md.push_str(&format!("- {}\n", note));
    }

    md
}

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

    // Query GPU info
    let gpu_name = query_gpu_name()?;
    let total_vram_gb = query_total_vram_gb()?;

    println!("GPU: {}", gpu_name);
    println!("Total VRAM: {:.1} GB\n", total_vram_gb);

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

    let mut results = Vec::new();
    let mut notes = Vec::new();

    // Test TFT with batch sizes: 16, 32, 64, 128
    println!("Testing TFT model...");
    for batch_size in [16, 32, 64, 128] {
        match benchmark_tft_batch(batch_size, &device) {
            Ok(result) => {
                println!(
                    "    Batch {}: VRAM={:.0}MB ({:.1}%), Throughput={:.1}/sec, Status={}",
                    result.batch_size,
                    result.vram_used_mb,
                    result.vram_percent,
                    result.throughput_samples_per_sec,
                    if result.oom_occurred { "OOM" } else { "OK" }
                );
                results.push(result);
            },
            Err(e) => {
                eprintln!("    Error: {}", e);
            },
        }
    }

    // Test MAMBA-2 with batch sizes: 8, 16, 32
    println!("\nTesting MAMBA-2 model...");
    for batch_size in [8, 16, 32] {
        match benchmark_mamba_batch(batch_size, &device) {
            Ok(result) => {
                println!(
                    "    Batch {}: VRAM={:.0}MB ({:.1}%), Throughput={:.1}/sec, Status={}",
                    result.batch_size,
                    result.vram_used_mb,
                    result.vram_percent,
                    result.throughput_samples_per_sec,
                    if result.oom_occurred { "OOM" } else { "OK" }
                );
                results.push(result);
            },
            Err(e) => {
                eprintln!("    Error: {}", e);
            },
        }
    }

    // Test Liquid with batch sizes: 16, 32, 64
    println!("\nTesting Liquid model (CPU)...");
    for batch_size in [16, 32, 64] {
        match benchmark_liquid_batch(batch_size) {
            Ok(result) => {
                println!(
                    "    Batch {}: Throughput={:.1}/sec, Status={}",
                    result.batch_size,
                    result.throughput_samples_per_sec,
                    if result.oom_occurred { "OOM" } else { "OK" }
                );
                results.push(result);
            },
            Err(e) => {
                eprintln!("    Error: {}", e);
            },
        }
    }

    // Determine optimal batch sizes
    let mut recommendations = HashMap::new();

    // TFT: Find largest batch size with <90% VRAM and no OOM
    if let Some(best) = results
        .iter()
        .filter(|r| r.model_name == "TFT" && !r.oom_occurred && r.vram_percent < 90.0)
        .max_by_key(|r| r.batch_size)
    {
        recommendations.insert("TFT".to_string(), best.batch_size);
    }

    // MAMBA-2: Find largest batch size with <90% VRAM and no OOM
    if let Some(best) = results
        .iter()
        .filter(|r| r.model_name == "MAMBA-2" && !r.oom_occurred && r.vram_percent < 90.0)
        .max_by_key(|r| r.batch_size)
    {
        recommendations.insert("MAMBA-2".to_string(), best.batch_size);
    }

    // Liquid: Find highest throughput (CPU-based)
    if let Some(best) = results
        .iter()
        .filter(|r| r.model_name == "Liquid" && !r.oom_occurred)
        .max_by(|a, b| {
            a.throughput_samples_per_sec
                .partial_cmp(&b.throughput_samples_per_sec)
                .unwrap()
        })
    {
        recommendations.insert("Liquid".to_string(), best.batch_size);
    }

    // Add notes
    notes.push(format!(
        "Tested on {} with {:.1}GB VRAM",
        gpu_name, total_vram_gb
    ));
    notes.push("Batch sizes optimized for <90% VRAM usage".to_string());
    notes.push("Liquid model runs on CPU (no GPU VRAM usage)".to_string());
    notes.push(format!(
        "Warmup iterations: {}, Benchmark iterations: {}",
        WARMUP_ITERATIONS, BENCHMARK_ITERATIONS
    ));

    let report = OptimizationReport {
        timestamp: chrono::Utc::now().to_rfc3339(),
        gpu_name,
        vram_total_gb: total_vram_gb,
        results,
        recommendations: recommendations.clone(),
        notes,
    };

    // Generate markdown report
    let md_content = generate_report(&report);

    // Write report
    let mut file = File::create("BATCH_SIZE_OPTIMIZATION_REPORT.md")?;
    file.write_all(md_content.as_bytes())?;

    println!("\n=== Optimization Complete ===\n");
    println!("Recommendations:");
    for (model, batch) in &recommendations {
        println!("  {} -> batch_size = {}", model, batch);
    }
    println!("\nReport written to: BATCH_SIZE_OPTIMIZATION_REPORT.md");

    Ok(())
}