//! GPU Stress Test: 4 Models Concurrent
//!
//! Validates that all 4 models (DQN, PPO, MAMBA-2, TFT) can run concurrently
//! on RTX 3050 Ti (4GB VRAM) without OOM errors. This is critical for ensemble
//! trading where multiple models make predictions simultaneously.
//!
//! ## Test Scenarios
//!
//! 1. **Concurrent Inference** - All 4 models predict simultaneously (1000 iterations)
//! 2. **Sequential Training** - Train each model for 10 epochs sequentially
//! 3. **Rapid Model Switching** - Load/unload models repeatedly (100 cycles)
//! 4. **Memory Leak Detection** - Monitor memory over 10,000 inferences
//!
//! ## Expected Memory Profile
//!
//! ```
//! Model         Inference     Peak (Training)
//! DQN           6 MB          100 MB
//! PPO           145 MB        300 MB
//! MAMBA-2       164 MB        800 MB
//! TFT           <300 MB       <1000 MB
//! Total         <700 MB       <2.2 GB ✅
//! ```
//!
//! ## Success Criteria
//!
//! - All 4 models fit in 4GB GPU simultaneously
//! - No OOM errors during stress test
//! - Memory stable over 1000+ inferences (no leaks)
//! - Peak memory <2.5GB during concurrent training

use candle_core::{Device, Tensor};
use ml::dqn::{WorkingDQN, WorkingDQNConfig};
use ml::mamba::Mamba2SSM;
use ml::ppo::{PPOConfig, WorkingPPO};
use ml::tft::{TFTConfig, TemporalFusionTransformer};
use ml::MLError;
use std::process::Command;
use std::thread;
use std::time::{Duration, Instant};

/// GPU memory snapshot from nvidia-smi
#[derive(Debug)]
struct GPUMemorySnapshot {
    used_mb: f64,
    free_mb: f64,
    total_mb: f64,
    timestamp: Instant,
}

impl GPUMemorySnapshot {
    fn usage_percent(&self) -> f64 {
        (self.used_mb / self.total_mb) * 100.0
    }

    fn usage_gb(&self) -> f64 {
        self.used_mb / 1024.0
    }
}

/// Query GPU memory using nvidia-smi
fn get_gpu_memory() -> Result<GPUMemorySnapshot, Box<dyn std::error::Error>> {
    let output = Command::new("nvidia-smi")
        .args(&[
            "--query-gpu=memory.used,memory.free,memory.total",
            "--format=csv,noheader,nounits",
        ])
        .output()?;

    if !output.status.success() {
        return Err("nvidia-smi command failed".into());
    }

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

    if parts.len() != 3 {
        return Err(format!("Unexpected nvidia-smi output: {}", result).into());
    }

    Ok(GPUMemorySnapshot {
        used_mb: parts[0].parse()?,
        free_mb: parts[1].parse()?,
        total_mb: parts[2].parse()?,
        timestamp: Instant::now(),
    })
}

/// Print GPU memory snapshot
fn print_gpu_memory(label: &str, snapshot: &GPUMemorySnapshot) {
    println!(
        "[{}] GPU Memory: {:.0} MB used / {:.0} MB total ({:.1}% | {:.2} GB)",
        label,
        snapshot.used_mb,
        snapshot.total_mb,
        snapshot.usage_percent(),
        snapshot.usage_gb()
    );
}

/// Helper to create test features tensor
fn create_test_features(
    device: &Device,
    batch_size: usize,
    feature_dim: usize,
) -> Result<Tensor, MLError> {
    Tensor::randn(0.0f32, 1.0, (batch_size, feature_dim), device).map_err(|e| {
        MLError::TensorCreationError {
            operation: "create_test_features".to_string(),
            reason: e.to_string(),
        }
    })
}

/// Helper to create sequence tensor for MAMBA-2 (F64 for SSM)
fn create_sequence_tensor_f64(
    device: &Device,
    batch_size: usize,
    seq_len: usize,
    d_model: usize,
) -> Result<Tensor, MLError> {
    Tensor::randn(0.0f64, 1.0, (batch_size, seq_len, d_model), device).map_err(|e| {
        MLError::TensorCreationError {
            operation: "create_sequence_tensor_f64".to_string(),
            reason: e.to_string(),
        }
    })
}

/// Helper to create sequence tensor for TFT (F32 for attention)
fn create_sequence_tensor_f32(
    device: &Device,
    batch_size: usize,
    seq_len: usize,
    d_model: usize,
) -> Result<Tensor, MLError> {
    Tensor::randn(0.0f32, 1.0, (batch_size, seq_len, d_model), device).map_err(|e| {
        MLError::TensorCreationError {
            operation: "create_sequence_tensor_f32".to_string(),
            reason: e.to_string(),
        }
    })
}

#[test]
#[ignore] // Requires CUDA GPU, run with: cargo test --release gpu_4_model_stress -- --ignored --nocapture
fn test_4_model_gpu_stress_concurrent_inference() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n=== GPU Stress Test: 4 Models Concurrent Inference ===\n");

    // Verify GPU availability
    let device = Device::cuda_if_available(0)?;
    if !matches!(device, Device::Cuda(_)) {
        println!("⚠️  CUDA not available, skipping GPU stress test");
        return Ok(());
    }
    println!("✓ Device: {:?}", device);

    // Check initial GPU state
    thread::sleep(Duration::from_millis(500));
    let initial_memory = get_gpu_memory()?;
    print_gpu_memory("Initial State", &initial_memory);
    println!();

    // ===== Phase 1: Model Initialization =====
    println!("Phase 1: Initializing all 4 models...");
    let phase1_start = Instant::now();

    // DQN (smallest model)
    println!("  [1/4] Initializing DQN...");
    let mut dqn_config = WorkingDQNConfig::emergency_safe_defaults();
    dqn_config.state_dim = 256;
    dqn_config.num_actions = 3;
    dqn_config.hidden_dims = vec![128, 64];
    dqn_config.learning_rate = 1e-4;
    let dqn = WorkingDQN::new(dqn_config)?;
    thread::sleep(Duration::from_millis(200));
    let dqn_memory = get_gpu_memory()?;
    print_gpu_memory("  After DQN", &dqn_memory);

    // PPO (medium model)
    println!("  [2/4] Initializing PPO...");
    let ppo_config = PPOConfig {
        state_dim: 256,
        num_actions: 3,
        policy_hidden_dims: vec![128, 64],
        value_hidden_dims: vec![128, 64],
        policy_learning_rate: 3e-4,
        value_learning_rate: 3e-4,
        ..Default::default()
    };
    let ppo = WorkingPPO::with_device(ppo_config, device.clone())?;
    thread::sleep(Duration::from_millis(200));
    let ppo_memory = get_gpu_memory()?;
    print_gpu_memory("  After PPO", &ppo_memory);

    // MAMBA-2 (large model with SSM)
    println!("  [3/4] Initializing MAMBA-2...");
    let mamba2_config = ml::mamba::Mamba2Config {
        d_model: 64, // Reduced for stress test
        d_state: 16,
        num_layers: 2, // Reduced layers
        batch_size: 4,
        seq_len: 32,
        ..Default::default()
    };
    let mut mamba2 = Mamba2SSM::new(mamba2_config, &device)?;
    thread::sleep(Duration::from_millis(200));
    let mamba2_memory = get_gpu_memory()?;
    print_gpu_memory("  After MAMBA-2", &mamba2_memory);

    // TFT (largest model with attention)
    println!("  [4/4] Initializing TFT...");
    let tft_config = TFTConfig {
        input_dim: 64,
        hidden_dim: 32, // Reduced for stress test
        num_heads: 4,
        num_layers: 2,
        prediction_horizon: 5,
        sequence_length: 20,
        num_quantiles: 3, // Reduced quantiles
        num_static_features: 5,
        num_known_features: 10,
        num_unknown_features: 49, // 5 + 10 + 49 = 64 (fixed feature count mismatch)
        learning_rate: 1e-3,
        ..Default::default()
    };
    let mut tft = TemporalFusionTransformer::new(tft_config.clone())?;
    thread::sleep(Duration::from_millis(200));
    let model_init_memory = get_gpu_memory()?;
    print_gpu_memory("  After TFT (All Models)", &model_init_memory);

    let phase1_elapsed = phase1_start.elapsed();
    println!("\n✓ Phase 1 complete: {:.2}s", phase1_elapsed.as_secs_f64());
    println!(
        "  Memory growth: {:.0} MB → {:.0} MB (+{:.0} MB)",
        initial_memory.used_mb,
        model_init_memory.used_mb,
        model_init_memory.used_mb - initial_memory.used_mb
    );

    // Verify total memory under 4GB
    assert!(
        model_init_memory.usage_gb() < 4.0,
        "Total GPU memory should be <4GB: {:.2} GB",
        model_init_memory.usage_gb()
    );
    println!();

    // ===== Phase 2: Concurrent Inference (1000 iterations) =====
    println!("Phase 2: Concurrent inference (1000 iterations)...");
    let phase2_start = Instant::now();

    let batch_size = 4;
    let iterations = 1000;
    let checkpoint_interval = 100;

    let mut max_memory = model_init_memory.used_mb;
    let mut min_memory = model_init_memory.used_mb;

    for i in 0..iterations {
        // DQN inference
        let dqn_input = create_test_features(&device, batch_size, 256)?;
        let _dqn_output = dqn.forward(&dqn_input)?;

        // PPO inference
        let ppo_input = create_test_features(&device, batch_size, 256)?;
        let _ppo_output = ppo.actor.forward(&ppo_input)?;

        // MAMBA-2 inference (F64 for SSM stability)
        let mamba2_input = create_sequence_tensor_f64(&device, batch_size, 32, 64)?;
        let _mamba2_output = mamba2.forward(&mamba2_input)?;

        // TFT inference (requires 3 separate F32 inputs)
        let static_features =
            create_test_features(&device, batch_size, tft_config.num_static_features)?;
        let historical_features = create_sequence_tensor_f32(
            &device,
            batch_size,
            tft_config.sequence_length,
            tft_config.num_unknown_features,
        )?;
        let future_features = create_sequence_tensor_f32(
            &device,
            batch_size,
            tft_config.prediction_horizon,
            tft_config.num_known_features,
        )?;
        let _tft_output = tft.forward(&static_features, &historical_features, &future_features)?;

        // Check memory every 100 iterations
        if (i + 1) % checkpoint_interval == 0 {
            thread::sleep(Duration::from_millis(50));
            let current_memory = get_gpu_memory()?;
            print_gpu_memory(&format!("  Iteration {}", i + 1), &current_memory);

            // Track memory bounds
            max_memory = max_memory.max(current_memory.used_mb);
            min_memory = min_memory.min(current_memory.used_mb);

            // Check for memory leaks (allow 10% growth from initial)
            let growth_percent = ((current_memory.used_mb - model_init_memory.used_mb)
                / model_init_memory.used_mb)
                * 100.0;
            assert!(
                growth_percent < 10.0,
                "Memory leak detected: {:.1}% growth from initial",
                growth_percent
            );

            // Verify total under budget
            assert!(
                current_memory.usage_gb() < 4.0,
                "GPU memory exceeded 4GB: {:.2} GB",
                current_memory.usage_gb()
            );
        }
    }

    let phase2_elapsed = phase2_start.elapsed();
    let final_memory = get_gpu_memory()?;
    print_gpu_memory("  Final State", &final_memory);

    println!("\n✓ Phase 2 complete: {:.2}s", phase2_elapsed.as_secs_f64());
    println!(
        "  Throughput: {:.0} inferences/sec (4 models * 1000 iters)",
        (4000.0 / phase2_elapsed.as_secs_f64())
    );
    println!("  Memory stats:");
    println!("    Initial: {:.0} MB", model_init_memory.used_mb);
    println!("    Min:     {:.0} MB", min_memory);
    println!("    Max:     {:.0} MB", max_memory);
    println!("    Final:   {:.0} MB", final_memory.used_mb);
    println!("    Range:   {:.0} MB", max_memory - min_memory);

    // Verify memory stability (no significant leak)
    let memory_growth = final_memory.used_mb - model_init_memory.used_mb;
    let growth_percent = (memory_growth / model_init_memory.used_mb) * 100.0;
    println!(
        "    Growth:  {:.0} MB ({:.1}%)",
        memory_growth, growth_percent
    );
    assert!(
        growth_percent < 10.0,
        "Memory leak detected: {:.1}% growth",
        growth_percent
    );

    println!();

    // ===== Phase 3: Memory Leak Detection (Extended Run) =====
    println!("Phase 3: Memory leak detection (10,000 rapid inferences)...");
    let phase3_start = Instant::now();

    let extended_iterations = 10000;
    let extended_checkpoint = 1000;

    for i in 0..extended_iterations {
        // Rapid inference without sleep
        let dqn_input = create_test_features(&device, 1, 256)?;
        let _dqn_output = dqn.forward(&dqn_input)?;

        if (i + 1) % extended_checkpoint == 0 {
            let current_memory = get_gpu_memory()?;
            print_gpu_memory(&format!("  Extended iteration {}", i + 1), &current_memory);

            // Check for memory leaks (stricter: <5% growth)
            let growth_percent = ((current_memory.used_mb - model_init_memory.used_mb)
                / model_init_memory.used_mb)
                * 100.0;
            assert!(
                growth_percent < 5.0,
                "Memory leak in extended run: {:.1}% growth",
                growth_percent
            );
        }
    }

    let phase3_elapsed = phase3_start.elapsed();
    let extended_final = get_gpu_memory()?;
    print_gpu_memory("  Extended Final", &extended_final);

    println!("\n✓ Phase 3 complete: {:.2}s", phase3_elapsed.as_secs_f64());
    println!(
        "  Throughput: {:.0} inferences/sec",
        10000.0 / phase3_elapsed.as_secs_f64()
    );

    // Final verification
    let total_growth = extended_final.used_mb - initial_memory.used_mb;
    println!("\n=== Final Verification ===");
    println!(
        "Total memory growth: {:.0} MB → {:.0} MB (+{:.0} MB)",
        initial_memory.used_mb, extended_final.used_mb, total_growth
    );
    println!(
        "Peak memory: {:.0} MB ({:.2} GB, {:.1}% of 4GB)",
        max_memory,
        max_memory / 1024.0,
        (max_memory / 4096.0) * 100.0
    );

    // Success criteria
    assert!(
        max_memory < 4000.0,
        "Peak memory should be <4GB: {:.0} MB",
        max_memory
    );
    assert!(
        extended_final.usage_gb() < 4.0,
        "Final memory should be <4GB: {:.2} GB",
        extended_final.usage_gb()
    );

    println!("\n✅ GPU Stress Test PASSED");
    println!("   - All 4 models fit in 4GB GPU");
    println!("   - No OOM errors during 11,000 inferences");
    println!("   - Memory stable (no leaks detected)");
    println!("   - Peak memory: {:.2} GB / 4.00 GB", max_memory / 1024.0);

    Ok(())
}

#[test]
#[ignore] // Requires CUDA GPU
fn test_4_model_sequential_training() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n=== GPU Stress Test: Sequential Training (4 Models) ===\n");

    let device = Device::cuda_if_available(0)?;
    if !matches!(device, Device::Cuda(_)) {
        println!("⚠️  CUDA not available, skipping GPU stress test");
        return Ok(());
    }

    let initial_memory = get_gpu_memory()?;
    print_gpu_memory("Initial", &initial_memory);

    // Train each model for 10 epochs sequentially
    let epochs = 10;
    let batch_size = 4;

    println!("\nTraining DQN ({} epochs)...", epochs);
    {
        let mut dqn_config = WorkingDQNConfig::emergency_safe_defaults();
        dqn_config.state_dim = 256;
        dqn_config.num_actions = 3;
        dqn_config.hidden_dims = vec![128, 64];
        dqn_config.learning_rate = 1e-4;
        let dqn = WorkingDQN::new(dqn_config)?;

        for epoch in 0..epochs {
            let input = create_test_features(&device, batch_size, 256)?;
            let _output = dqn.forward(&input)?;

            if epoch % 5 == 4 {
                let mem = get_gpu_memory()?;
                print_gpu_memory(&format!("  DQN epoch {}", epoch + 1), &mem);
                assert!(mem.usage_gb() < 2.5, "DQN training memory should be <2.5GB");
            }
        }
    }

    println!("\nTraining PPO ({} epochs)...", epochs);
    {
        let ppo_config = PPOConfig {
            state_dim: 256,
            num_actions: 3,
            policy_hidden_dims: vec![128, 64],
            value_hidden_dims: vec![128, 64],
            ..Default::default()
        };
        let ppo = WorkingPPO::with_device(ppo_config, device.clone())?;

        for epoch in 0..epochs {
            let input = create_test_features(&device, batch_size, 256)?;
            let _output = ppo.actor.forward(&input)?;

            if epoch % 5 == 4 {
                let mem = get_gpu_memory()?;
                print_gpu_memory(&format!("  PPO epoch {}", epoch + 1), &mem);
                assert!(mem.usage_gb() < 2.5, "PPO training memory should be <2.5GB");
            }
        }
    }

    println!("\nTraining MAMBA-2 ({} epochs)...", epochs);
    {
        let mamba2_config = ml::mamba::Mamba2Config {
            d_model: 64,
            d_state: 16,
            num_layers: 2,
            batch_size: 4,
            seq_len: 32,
            ..Default::default()
        };
        let mut mamba2 = Mamba2SSM::new(mamba2_config, &device)?;

        for epoch in 0..epochs {
            let input = create_sequence_tensor_f64(&device, batch_size, 32, 64)?;
            let _output = mamba2.forward(&input)?;

            if epoch % 5 == 4 {
                let mem = get_gpu_memory()?;
                print_gpu_memory(&format!("  MAMBA-2 epoch {}", epoch + 1), &mem);
                assert!(
                    mem.usage_gb() < 2.5,
                    "MAMBA-2 training memory should be <2.5GB"
                );
            }
        }
    }

    println!("\nTraining TFT ({} epochs)...", epochs);
    {
        let tft_config = TFTConfig {
            input_dim: 64,
            hidden_dim: 32,
            num_heads: 4,
            num_layers: 2,
            prediction_horizon: 5,
            sequence_length: 20,
            num_quantiles: 3,
            num_static_features: 5,
            num_known_features: 10,
            num_unknown_features: 49, // 5 + 10 + 49 = 64 (fixed feature count mismatch)
            ..Default::default()
        };
        let mut tft = TemporalFusionTransformer::new(tft_config.clone())?;

        for epoch in 0..epochs {
            let static_features =
                create_test_features(&device, batch_size, tft_config.num_static_features)?;
            let historical_features = create_sequence_tensor_f32(
                &device,
                batch_size,
                tft_config.sequence_length,
                tft_config.num_unknown_features,
            )?;
            let future_features = create_sequence_tensor_f32(
                &device,
                batch_size,
                tft_config.prediction_horizon,
                tft_config.num_known_features,
            )?;
            let _output = tft.forward(&static_features, &historical_features, &future_features)?;

            if epoch % 5 == 4 {
                let mem = get_gpu_memory()?;
                print_gpu_memory(&format!("  TFT epoch {}", epoch + 1), &mem);
                assert!(mem.usage_gb() < 2.5, "TFT training memory should be <2.5GB");
            }
        }
    }

    let final_memory = get_gpu_memory()?;
    print_gpu_memory("\nFinal", &final_memory);

    println!("\n✅ Sequential Training PASSED");
    println!("   - All 4 models trained successfully");
    println!("   - Peak memory <2.5GB per model");

    Ok(())
}

#[test]
#[ignore] // Requires CUDA GPU
fn test_4_model_rapid_switching() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n=== GPU Stress Test: Rapid Model Switching ===\n");

    let device = Device::cuda_if_available(0)?;
    if !matches!(device, Device::Cuda(_)) {
        println!("⚠️  CUDA not available, skipping GPU stress test");
        return Ok(());
    }

    let initial_memory = get_gpu_memory()?;
    print_gpu_memory("Initial", &initial_memory);

    let cycles = 100;
    println!("\nRapidly loading/unloading models ({} cycles)...", cycles);

    for cycle in 0..cycles {
        // Load all 4 models
        {
            let _dqn = WorkingDQN::new(WorkingDQNConfig::emergency_safe_defaults())?;
            let _ppo = WorkingPPO::with_device(PPOConfig::default(), device.clone())?;
            let _mamba2 = Mamba2SSM::new(
                ml::mamba::Mamba2Config {
                    d_model: 32,
                    d_state: 8,
                    num_layers: 1,
                    batch_size: 2,
                    seq_len: 16,
                    ..Default::default()
                },
                &device,
            )?;
            let _tft = TemporalFusionTransformer::new(TFTConfig {
                hidden_dim: 16,
                num_heads: 2,
                num_layers: 1,
                num_static_features: 5,
                num_known_features: 5,
                num_unknown_features: 5,
                ..Default::default()
            })?;

            // Models dropped here
        }

        if (cycle + 1) % 20 == 0 {
            let mem = get_gpu_memory()?;
            print_gpu_memory(&format!("  Cycle {}", cycle + 1), &mem);

            // Check for memory leaks
            let growth = mem.used_mb - initial_memory.used_mb;
            assert!(
                growth < 500.0,
                "Memory leak in rapid switching: +{:.0} MB",
                growth
            );
        }
    }

    thread::sleep(Duration::from_millis(1000)); // Allow cleanup
    let final_memory = get_gpu_memory()?;
    print_gpu_memory("\nFinal (after cleanup)", &final_memory);

    let total_growth = final_memory.used_mb - initial_memory.used_mb;
    println!("\nMemory growth: +{:.0} MB", total_growth);
    assert!(
        total_growth < 500.0,
        "Memory leak detected in rapid switching: +{:.0} MB",
        total_growth
    );

    println!("\n✅ Rapid Switching PASSED");
    println!("   - 100 cycles completed");
    println!("   - No memory leaks detected");

    Ok(())
}