foxhunt/ml/tests/inference_optimization_tests.rs

//! ML Inference Optimization Tests
//!
//! Comprehensive testing for ML inference optimizations including:
//! - Batch inference with dynamic batching
//! - Model quantization (INT8, FP16) with accuracy measurement
//! - Model pruning (structured/unstructured) with sparsity analysis
//! - GPU vs CPU inference latency comparison
//! - Concurrent inference requests (thread safety)
//! - Memory usage per inference
//! - Model warmup effects
//! - Variable batch sizes with padding/masking
//!
//! **Performance Targets** (from HFT requirements):
//! - Single inference: <50μs (latency-critical, batch=1)
//! - Batch inference: >1000 inferences/sec (throughput optimization)
//! - GPU speedup: 10-50x vs CPU for large models
//! - Quantization accuracy loss: <2% for FP16, <5% for INT8
//! - Memory per inference: <10MB for production models

#![allow(unused_crate_dependencies)]

use std::sync::Arc;
use std::time::{Duration, Instant};

use candle_core::{DType, Device, Tensor};
use ml::features::FeatureVector;
use ml::inference::{ModelConfig, RealInferenceConfig, RealMLInferenceEngine, RealNeuralNetwork};
use ml::safety::{MLSafetyConfig, MLSafetyManager};

// Helper function to create mock features for testing
// Returns a 225-dimension feature vector filled with normalized test data
fn create_mock_features() -> FeatureVector {
    let mut features = [0.0f64; 225];

    // Fill with realistic test data (normalized values between -3 and 3 for z-score)
    for (i, val) in features.iter_mut().enumerate() {
        *val = ((i as f64) / 225.0) * 6.0 - 3.0; // Range: -3.0 to 3.0
    }

    features
}

// ==================== LATENCY TESTS ====================

/// Test: Batch inference with size 1 (latency-critical path)
#[tokio::test]
async fn test_batch_inference_size_1_latency() -> Result<(), Box<dyn std::error::Error>> {
    let safety_manager = Arc::new(MLSafetyManager::new(MLSafetyConfig::default()));
    let mut config = RealInferenceConfig::default();
    config.batch_size = 1;
    config.device_preference = "cpu".to_string();
    config.enable_caching = false; // Disable cache to measure pure inference

    let engine = RealMLInferenceEngine::new(config, safety_manager);

    let model_config = ModelConfig {
        input_dim: 21,
        hidden_dims: vec![32, 16],
        output_dim: 1,
        activation: "relu".to_string(),
        batch_norm: false,
        dropout_rate: 0.0,
    };

    engine
        .load_model("latency_test".to_string(), model_config)
        .await?;

    let features = create_mock_features();

    // Warmup (first inference is always slower)
    let _ = engine.predict("latency_test", &features).await;

    // Measure inference latency (should be <50μs for HFT)
    let start = Instant::now();
    let result = engine.predict("latency_test", &features).await?;
    let latency_us = start.elapsed().as_micros();

    println!("Single inference latency: {}μs", latency_us);

    // Verify result structure
    assert!(result.confidence > 0.0);
    assert!(result.inference_latency_us > 0);

    // HFT target: <50μs (might not hit on CPU, but should be reasonable)
    // Note: This test just measures, doesn't enforce strict limit
    assert!(latency_us < 10_000); // Sanity check: <10ms

    Ok(())
}

/// Test: Batch size optimization - throughput vs latency trade-off
#[tokio::test]
async fn test_batch_size_throughput_optimization() -> Result<(), Box<dyn std::error::Error>> {
    let safety_manager = Arc::new(MLSafetyManager::new(MLSafetyConfig::default()));
    let mut config = RealInferenceConfig::default();
    config.device_preference = "cpu".to_string();
    config.enable_caching = false;

    let batch_sizes = vec![1, 2, 4, 8];
    let mut throughputs = Vec::new();

    for batch_size in &batch_sizes {
        config.batch_size = *batch_size;
        let engine = RealMLInferenceEngine::new(config.clone(), Arc::clone(&safety_manager));

        let model_config = ModelConfig {
            input_dim: 21,
            hidden_dims: vec![32],
            output_dim: 1,
            activation: "relu".to_string(),
            batch_norm: false,
            dropout_rate: 0.0,
        };

        engine
            .load_model("throughput_test".to_string(), model_config)
            .await?;

        // Warmup
        let features = create_mock_features();
        let _ = engine.predict("throughput_test", &features).await;

        // Measure throughput: inferences per second
        let num_inferences = 20;
        let start = Instant::now();

        for _ in 0..num_inferences {
            let features = create_mock_features();
            let _ = engine.predict("throughput_test", &features).await;
        }

        let elapsed_secs = start.elapsed().as_secs_f64();
        let throughput = num_inferences as f64 / elapsed_secs;
        throughputs.push(throughput);

        println!(
            "Batch size {}: {:.2} inferences/sec",
            batch_size, throughput
        );
    }

    // Verify throughput increases with batch size (or stays reasonable)
    // Larger batches should generally be more efficient
    assert!(throughputs[0] > 0.0);
    assert!(throughputs.last().copied().unwrap_or(0.0) > 0.0);

    Ok(())
}

/// Test: Variable batch sizes with padding/masking
#[tokio::test]
async fn test_variable_batch_sizes() -> Result<(), Box<dyn std::error::Error>> {
    let device = Device::Cpu;

    let model_config = ModelConfig {
        input_dim: 10,
        hidden_dims: vec![20],
        output_dim: 1,
        activation: "relu".to_string(),
        batch_norm: false,
        dropout_rate: 0.0,
    };

    let network = RealNeuralNetwork::new(model_config, device.clone())?;

    // Test different batch sizes (1, 2, 4, 8)
    let batch_sizes = vec![1, 2, 4, 8];

    for batch_size in batch_sizes {
        let input_data = vec![1.0f32; batch_size * 10];
        let input = Tensor::from_vec(input_data, &[batch_size, 10], &device)?;

        let output = network.forward(&input).await?;
        let output_dims = output.dims();

        // Verify output shape matches batch size
        assert_eq!(output_dims.len(), 2);
        assert_eq!(output_dims[0], batch_size);
        assert_eq!(output_dims[1], 1);

        println!(
            "Batch size {} -> output shape: {:?}",
            batch_size, output_dims
        );
    }

    Ok(())
}

/// Test: Large batch size - throughput optimization
#[tokio::test]
async fn test_large_batch_inference() -> Result<(), Box<dyn std::error::Error>> {
    let device = Device::Cpu;

    let model_config = ModelConfig {
        input_dim: 21,
        hidden_dims: vec![32],
        output_dim: 1,
        activation: "relu".to_string(),
        batch_norm: false,
        dropout_rate: 0.0,
    };

    let network = RealNeuralNetwork::new(model_config, device.clone())?;

    // Large batch for throughput (32 samples)
    let batch_size = 32;
    let input_data = vec![1.0f32; batch_size * 21];
    let input = Tensor::from_vec(input_data, &[batch_size, 21], &device)?;

    let start = Instant::now();
    let output = network.forward(&input).await?;
    let latency_us = start.elapsed().as_micros();

    let output_dims = output.dims();
    assert_eq!(output_dims[0], batch_size);

    // Calculate per-sample latency
    let per_sample_us = latency_us as f64 / batch_size as f64;
    println!(
        "Batch {} inference: {}μs total, {:.2}μs per sample",
        batch_size, latency_us, per_sample_us
    );

    // Batch inference should amortize cost (per-sample latency < single inference)
    assert!(per_sample_us < 1000.0); // Sanity check: <1ms per sample

    Ok(())
}

// ==================== QUANTIZATION TESTS ====================

/// Test: FP32 vs FP16 quantization - accuracy comparison
#[tokio::test]
async fn test_quantization_fp16_accuracy() -> Result<(), Box<dyn std::error::Error>> {
    let device = Device::Cpu;

    let model_config = ModelConfig {
        input_dim: 10,
        hidden_dims: vec![20, 10],
        output_dim: 1,
        activation: "relu".to_string(),
        batch_norm: false,
        dropout_rate: 0.0,
    };

    let network = RealNeuralNetwork::new(model_config, device.clone())?;

    // Create test input
    let input_data = vec![1.0f32; 10];
    let input_fp32 = Tensor::from_vec(input_data.clone(), &[1, 10], &device)?;

    // FP32 inference (baseline)
    let output_fp32 = network.forward(&input_fp32).await?;
    let fp32_value = output_fp32
        .get(0)?
        .get(0)?
        .to_scalar::<f32>()
        .map_err(|e| format!("Failed to extract FP32 scalar: {}", e))?;

    // FP16 inference: Can't use F16 dtype directly with existing weights
    // Instead we simulate FP16 quantization effects by:
    // 1. Converting to F16 (lossy conversion)
    // 2. Converting back to F32 for computation
    // This demonstrates the accuracy impact without mixed-dtype matmul
    let input_fp16_sim = input_fp32.to_dtype(DType::F16)?.to_dtype(DType::F32)?;
    let output_fp16 = network.forward(&input_fp16_sim).await?;
    let fp16_value = output_fp16
        .get(0)?
        .get(0)?
        .to_scalar::<f32>()
        .map_err(|e| format!("Failed to extract FP16 scalar: {}", e))?;

    // Calculate accuracy degradation
    let diff = (fp32_value - fp16_value).abs();
    let relative_error = if fp32_value.abs() > 1e-6 {
        diff / fp32_value.abs()
    } else {
        diff
    };

    println!(
        "FP32: {:.6}, FP16: {:.6}, Relative Error: {:.6} ({:.2}%)",
        fp32_value,
        fp16_value,
        relative_error,
        relative_error * 100.0
    );

    // FP16 accuracy target: <2% error
    // Note: With random weights, error might be higher, but should be finite
    assert!(relative_error.is_finite());
    assert!(relative_error < 0.5); // Relaxed for random weights

    Ok(())
}

/// Test: INT8 quantization simulation - accuracy loss measurement
#[tokio::test]
async fn test_quantization_int8_simulation() -> Result<(), Box<dyn std::error::Error>> {
    let device = Device::Cpu;

    let model_config = ModelConfig {
        input_dim: 10,
        hidden_dims: vec![20],
        output_dim: 1,
        activation: "tanh".to_string(), // Tanh for bounded output
        batch_norm: false,
        dropout_rate: 0.0,
    };

    let network = RealNeuralNetwork::new(model_config, device.clone())?;

    // Create test input
    let input_data = vec![0.5f32; 10];
    let input = Tensor::from_vec(input_data, &[1, 10], &device)?;

    // FP32 inference (baseline)
    let output_fp32 = network.forward(&input).await?;
    let fp32_value = output_fp32
        .get(0)?
        .get(0)?
        .to_scalar::<f32>()
        .map_err(|e| format!("Failed to extract FP32 scalar: {}", e))?;

    // Simulate INT8 quantization: scale to [-127, 127], round, scale back
    let scale_factor = 127.0;
    let input_int8_sim = (input.clone() * scale_factor)?;
    let input_int8_sim = input_int8_sim.round()?;
    let input_int8_sim = (input_int8_sim / scale_factor)?;

    let output_int8 = network.forward(&input_int8_sim).await?;
    let int8_value = output_int8
        .get(0)?
        .get(0)?
        .to_scalar::<f32>()
        .map_err(|e| format!("Failed to extract INT8 scalar: {}", e))?;

    // Calculate accuracy degradation
    let diff = (fp32_value - int8_value).abs();
    let relative_error = if fp32_value.abs() > 1e-6 {
        diff / fp32_value.abs()
    } else {
        diff
    };

    println!(
        "FP32: {:.6}, INT8: {:.6}, Relative Error: {:.6} ({:.2}%)",
        fp32_value,
        int8_value,
        relative_error,
        relative_error * 100.0
    );

    // INT8 quantization with random untrained weights can have very high error
    // The key test is that values remain finite (no NaN/Inf)
    // In production with trained models, error would be much lower (<5%)
    assert!(relative_error.is_finite());
    assert!(fp32_value.is_finite());
    assert!(int8_value.is_finite());

    Ok(())
}

/// Test: Quantization memory savings
#[test]
fn test_quantization_memory_savings() {
    let num_params = 1000; // 1K parameters

    // FP32: 4 bytes per parameter
    let fp32_bytes = num_params * 4;

    // FP16: 2 bytes per parameter
    let fp16_bytes = num_params * 2;

    // INT8: 1 byte per parameter
    let int8_bytes = num_params * 1;

    let fp16_saving = (fp32_bytes - fp16_bytes) as f64 / fp32_bytes as f64;
    let int8_saving = (fp32_bytes - int8_bytes) as f64 / fp32_bytes as f64;

    println!("FP32: {} bytes", fp32_bytes);
    println!(
        "FP16: {} bytes ({:.1}% saving)",
        fp16_bytes,
        fp16_saving * 100.0
    );
    println!(
        "INT8: {} bytes ({:.1}% saving)",
        int8_bytes,
        int8_saving * 100.0
    );

    // FP16 should save 50%
    assert!((fp16_saving - 0.5).abs() < 0.01);

    // INT8 should save 75%
    assert!((int8_saving - 0.75).abs() < 0.01);
}

// ==================== MODEL PRUNING TESTS ====================

/// Test: Structured pruning - remove entire channels
#[tokio::test]
async fn test_structured_pruning_channels() -> Result<(), Box<dyn std::error::Error>> {
    let device = Device::Cpu;

    // Original model: 20 hidden units
    let config_original = ModelConfig {
        input_dim: 10,
        hidden_dims: vec![20],
        output_dim: 1,
        activation: "relu".to_string(),
        batch_norm: false,
        dropout_rate: 0.0,
    };

    // Pruned model: 10 hidden units (50% pruning)
    let config_pruned = ModelConfig {
        input_dim: 10,
        hidden_dims: vec![10],
        output_dim: 1,
        activation: "relu".to_string(),
        batch_norm: false,
        dropout_rate: 0.0,
    };

    let network_original = RealNeuralNetwork::new(config_original.clone(), device.clone())?;
    let network_pruned = RealNeuralNetwork::new(config_pruned.clone(), device.clone())?;

    // Measure inference latency
    let input_data = vec![1.0f32; 10];
    let input = Tensor::from_vec(input_data, &[1, 10], &device)?;

    let start = Instant::now();
    let _ = network_original.forward(&input).await?;
    let latency_original = start.elapsed().as_micros();

    let start = Instant::now();
    let _ = network_pruned.forward(&input).await?;
    let latency_pruned = start.elapsed().as_micros();

    let speedup = latency_original as f64 / (latency_pruned as f64).max(1.0);

    println!(
        "Original (20 units): {}μs, Pruned (10 units): {}μs, Speedup: {:.2}x",
        latency_original, latency_pruned, speedup
    );

    // Pruned model should be faster (or at least not slower)
    assert!(latency_pruned > 0);

    // Model sizes
    let original_params = 10 * 20 + 20 * 1; // input->hidden + hidden->output
    let pruned_params = 10 * 10 + 10 * 1;
    let size_reduction = (original_params - pruned_params) as f64 / original_params as f64;

    println!(
        "Parameter reduction: {:.1}% ({} -> {} params)",
        size_reduction * 100.0,
        original_params,
        pruned_params
    );

    assert!((size_reduction - 0.52).abs() < 0.1); // ~50% size reduction

    Ok(())
}

/// Test: Unstructured pruning - sparsity measurement
#[test]
fn test_unstructured_pruning_sparsity() {
    // Simulate weight pruning by zeroing out small weights
    let weights = vec![0.5, 0.01, 0.8, 0.02, 0.3, 0.001, 0.9, 0.05];
    let threshold = 0.1;

    let pruned_weights: Vec<f32> = weights
        .iter()
        .map(|&w: &f32| if w.abs() < threshold { 0.0 } else { w })
        .collect();

    let num_pruned = pruned_weights.iter().filter(|&&w| w == 0.0).count();
    let sparsity = num_pruned as f64 / pruned_weights.len() as f64;

    println!("Original weights: {:?}", weights);
    println!("Pruned weights: {:?}", pruned_weights);
    println!("Sparsity: {:.1}%", sparsity * 100.0);

    // Weights below threshold (0.1): 0.01, 0.02, 0.001, 0.05 = 4 out of 8
    // Expected sparsity: 4/8 = 50%
    assert!((sparsity - 0.5).abs() < 0.01);
}

/// Test: Pruning ratio vs accuracy trade-off
#[tokio::test]
async fn test_pruning_accuracy_tradeoff() -> Result<(), Box<dyn std::error::Error>> {
    let device = Device::Cpu;

    let hidden_sizes = vec![64, 48, 32, 16]; // 0%, 25%, 50%, 75% pruning

    for &hidden_size in &hidden_sizes {
        let config = ModelConfig {
            input_dim: 21,
            hidden_dims: vec![hidden_size],
            output_dim: 1,
            activation: "relu".to_string(),
            batch_norm: false,
            dropout_rate: 0.0,
        };

        let network = RealNeuralNetwork::new(config, device.clone())?;

        let input_data = vec![1.0f32; 21];
        let input = Tensor::from_vec(input_data, &[1, 21], &device)?;

        let start = Instant::now();
        let output = network.forward(&input).await?;
        let latency = start.elapsed().as_micros();

        let value = output
            .get(0)?
            .get(0)?
            .to_scalar::<f32>()
            .map_err(|e| format!("Failed to extract scalar: {}", e))?;

        let pruning_ratio = 1.0 - (hidden_size as f64 / 64.0);

        println!(
            "Hidden size: {}, Pruning: {:.1}%, Latency: {}μs, Output: {:.6}",
            hidden_size,
            pruning_ratio * 100.0,
            latency,
            value
        );
    }

    Ok(())
}

// ==================== GPU vs CPU INFERENCE TESTS ====================

/// Test: GPU vs CPU inference latency comparison (requires CUDA)
#[tokio::test]
#[ignore = "Requires CUDA GPU - skip in CI"]
async fn test_gpu_vs_cpu_latency() -> Result<(), Box<dyn std::error::Error>> {
    let model_config = ModelConfig {
        input_dim: 21,
        hidden_dims: vec![128, 64, 32],
        output_dim: 1,
        activation: "relu".to_string(),
        batch_norm: false,
        dropout_rate: 0.0,
    };

    // CPU inference
    let cpu_device = Device::Cpu;
    let network_cpu = RealNeuralNetwork::new(model_config.clone(), cpu_device.clone())?;

    let input_data = vec![1.0f32; 21];
    let input_cpu = Tensor::from_vec(input_data.clone(), &[1, 21], &cpu_device)?;

    // Warmup
    let _ = network_cpu.forward(&input_cpu).await?;

    let start = Instant::now();
    let _ = network_cpu.forward(&input_cpu).await?;
    let cpu_latency = start.elapsed().as_micros();

    // GPU inference
    let gpu_device = Device::new_cuda(0)?;
    let network_gpu = RealNeuralNetwork::new(model_config.clone(), gpu_device.clone())?;

    let input_gpu = Tensor::from_vec(input_data, &[1, 21], &gpu_device)?;

    // Warmup
    let _ = network_gpu.forward(&input_gpu).await?;

    let start = Instant::now();
    let _ = network_gpu.forward(&input_gpu).await?;
    let gpu_latency = start.elapsed().as_micros();

    let speedup = cpu_latency as f64 / (gpu_latency as f64).max(1.0);

    println!("CPU latency: {}μs", cpu_latency);
    println!("GPU latency: {}μs", gpu_latency);
    println!("GPU speedup: {:.2}x", speedup);

    // GPU should be faster for large models (or at least not catastrophically slower)
    // Note: Small models might not benefit from GPU due to overhead
    assert!(gpu_latency > 0);

    Ok(())
}

/// Test: GPU batch processing - large batch advantage
#[tokio::test]
#[ignore = "Requires CUDA GPU - skip in CI"]
async fn test_gpu_batch_advantage() -> Result<(), Box<dyn std::error::Error>> {
    let gpu_device = Device::new_cuda(0)?;

    let model_config = ModelConfig {
        input_dim: 21,
        hidden_dims: vec![128, 64],
        output_dim: 1,
        activation: "relu".to_string(),
        batch_norm: false,
        dropout_rate: 0.0,
    };

    let network = RealNeuralNetwork::new(model_config, gpu_device.clone())?;

    let batch_sizes = vec![1, 8, 32, 128];

    for batch_size in batch_sizes {
        let input_data = vec![1.0f32; batch_size * 21];
        let input = Tensor::from_vec(input_data, &[batch_size, 21], &gpu_device)?;

        // Warmup
        let _ = network.forward(&input).await?;

        let start = Instant::now();
        let _ = network.forward(&input).await?;
        let latency = start.elapsed().as_micros();

        let per_sample = latency as f64 / batch_size as f64;

        println!(
            "Batch {}: {}μs total, {:.2}μs per sample",
            batch_size, latency, per_sample
        );
    }

    Ok(())
}

// ==================== CONCURRENT INFERENCE TESTS ====================

/// Test: Concurrent inference requests - thread safety
#[tokio::test]
async fn test_concurrent_inference_thread_safety() -> Result<(), Box<dyn std::error::Error>> {
    let safety_manager = Arc::new(MLSafetyManager::new(MLSafetyConfig::default()));
    let mut config = RealInferenceConfig::default();
    config.device_preference = "cpu".to_string();

    let engine = Arc::new(RealMLInferenceEngine::new(config, safety_manager));

    let model_config = ModelConfig {
        input_dim: 21,
        hidden_dims: vec![32],
        output_dim: 1,
        activation: "relu".to_string(),
        batch_norm: false,
        dropout_rate: 0.0,
    };

    engine
        .load_model("concurrent_test".to_string(), model_config)
        .await?;

    // Spawn 10 concurrent inference tasks
    let num_concurrent = 10;
    let mut handles = Vec::new();

    for i in 0..num_concurrent {
        let engine_clone = Arc::clone(&engine);
        let handle = tokio::spawn(async move {
            let features = create_mock_features();
            let result = engine_clone.predict("concurrent_test", &features).await;
            (i, result)
        });
        handles.push(handle);
    }

    // Wait for all tasks and verify all succeeded
    let mut successes = 0;
    for handle in handles {
        let (task_id, result) = handle.await?;
        match result {
            Ok(_) => {
                successes += 1;
            },
            Err(e) => {
                println!("Task {} failed: {:?}", task_id, e);
            },
        }
    }

    println!(
        "Concurrent inferences: {}/{} succeeded",
        successes, num_concurrent
    );

    // All concurrent requests should succeed
    assert_eq!(successes, num_concurrent);

    Ok(())
}

/// Test: Concurrent inference with rate limiting
#[tokio::test]
async fn test_concurrent_inference_rate_limiting() -> Result<(), Box<dyn std::error::Error>> {
    let safety_manager = Arc::new(MLSafetyManager::new(MLSafetyConfig::default()));
    let mut config = RealInferenceConfig::default();
    config.device_preference = "cpu".to_string();
    config.max_inference_latency_us = 10_000; // 10ms timeout

    let engine = Arc::new(RealMLInferenceEngine::new(config, safety_manager));

    let model_config = ModelConfig {
        input_dim: 21,
        hidden_dims: vec![32],
        output_dim: 1,
        activation: "relu".to_string(),
        batch_norm: false,
        dropout_rate: 0.0,
    };

    engine
        .load_model("rate_test".to_string(), model_config)
        .await?;

    // Burst of requests
    let num_requests = 20;
    let start = Instant::now();

    let mut handles = Vec::new();
    for _ in 0..num_requests {
        let engine_clone = Arc::clone(&engine);
        let handle = tokio::spawn(async move {
            let features = create_mock_features();
            engine_clone.predict("rate_test", &features).await
        });
        handles.push(handle);
    }

    // Wait for all
    let mut successes = 0;
    for handle in handles {
        if handle.await?.is_ok() {
            successes += 1;
        }
    }

    let elapsed = start.elapsed();
    let qps = num_requests as f64 / elapsed.as_secs_f64();

    println!(
        "Processed {} requests in {:.3}s ({:.2} QPS)",
        successes,
        elapsed.as_secs_f64(),
        qps
    );

    assert!(successes > 0);

    Ok(())
}

// ==================== MEMORY USAGE TESTS ====================

/// Test: Memory usage per inference
#[tokio::test]
async fn test_memory_usage_per_inference() -> Result<(), Box<dyn std::error::Error>> {
    let safety_manager = Arc::new(MLSafetyManager::new(MLSafetyConfig::default()));
    let mut config = RealInferenceConfig::default();
    config.device_preference = "cpu".to_string();
    config.enable_caching = false;

    let engine = RealMLInferenceEngine::new(config, safety_manager);

    let model_config = ModelConfig {
        input_dim: 21,
        hidden_dims: vec![128, 64],
        output_dim: 1,
        activation: "relu".to_string(),
        batch_norm: false,
        dropout_rate: 0.0,
    };

    engine
        .load_model("memory_test".to_string(), model_config)
        .await?;

    let features = create_mock_features();
    let result = engine.predict("memory_test", &features).await?;

    let memory_bytes = result.memory_used_bytes;
    let memory_kb = memory_bytes as f64 / 1024.0;
    let memory_mb = memory_kb / 1024.0;

    println!(
        "Memory per inference: {} bytes ({:.2} KB, {:.2} MB)",
        memory_bytes, memory_kb, memory_mb
    );

    // Production target: <10MB per inference
    assert!(memory_bytes > 0);
    assert!(memory_mb < 50.0); // Sanity check: <50MB

    Ok(())
}

/// Test: Memory usage scaling with batch size
#[tokio::test]
async fn test_memory_scaling_batch_size() -> Result<(), Box<dyn std::error::Error>> {
    let device = Device::Cpu;

    let model_config = ModelConfig {
        input_dim: 21,
        hidden_dims: vec![64],
        output_dim: 1,
        activation: "relu".to_string(),
        batch_norm: false,
        dropout_rate: 0.0,
    };

    let network = RealNeuralNetwork::new(model_config, device.clone())?;

    let batch_sizes = vec![1, 4, 16, 64];

    for batch_size in batch_sizes {
        let input_data = vec![1.0f32; batch_size * 21];
        let input = Tensor::from_vec(input_data, &[batch_size, 21], &device)?;

        let _ = network.forward(&input).await?;

        // Estimate memory: batch_size * 21 features * 4 bytes (f32)
        let input_memory = batch_size * 21 * 4;
        let output_memory = batch_size * 1 * 4;
        let total_memory = input_memory + output_memory;

        println!(
            "Batch {}: ~{} bytes ({:.2} KB)",
            batch_size,
            total_memory,
            total_memory as f64 / 1024.0
        );

        assert!(total_memory > 0);
    }

    Ok(())
}

// ==================== MODEL WARMUP TESTS ====================

/// Test: First inference slower (warmup effect)
#[tokio::test]
async fn test_inference_warmup_effect() -> Result<(), Box<dyn std::error::Error>> {
    let safety_manager = Arc::new(MLSafetyManager::new(MLSafetyConfig::default()));
    let mut config = RealInferenceConfig::default();
    config.device_preference = "cpu".to_string();
    config.enable_caching = false;

    let engine = RealMLInferenceEngine::new(config, safety_manager);

    let model_config = ModelConfig {
        input_dim: 21,
        hidden_dims: vec![32],
        output_dim: 1,
        activation: "relu".to_string(),
        batch_norm: false,
        dropout_rate: 0.0,
    };

    engine
        .load_model("warmup_test".to_string(), model_config)
        .await?;

    let features = create_mock_features();

    // First inference (cold start)
    let start = Instant::now();
    let _ = engine.predict("warmup_test", &features).await?;
    let first_latency = start.elapsed().as_micros();

    // Subsequent inferences (warm)
    let num_warm = 5;
    let mut warm_latencies = Vec::new();

    for _ in 0..num_warm {
        let features = create_mock_features();
        let start = Instant::now();
        let _ = engine.predict("warmup_test", &features).await?;
        warm_latencies.push(start.elapsed().as_micros());
    }

    let avg_warm_latency = warm_latencies.iter().sum::<u128>() / num_warm as u128;

    println!("First inference (cold): {}μs", first_latency);
    println!("Average warm inference: {}μs", avg_warm_latency);
    println!(
        "Warmup overhead: {:.2}x slower",
        first_latency as f64 / avg_warm_latency as f64
    );

    // First inference might be slower (but not by orders of magnitude)
    assert!(first_latency > 0);
    assert!(avg_warm_latency > 0);

    Ok(())
}

/// Test: Model warmup iterations
#[tokio::test]
async fn test_model_warmup_iterations() -> Result<(), Box<dyn std::error::Error>> {
    let device = Device::Cpu;

    let model_config = ModelConfig {
        input_dim: 21,
        hidden_dims: vec![64, 32],
        output_dim: 1,
        activation: "relu".to_string(),
        batch_norm: false,
        dropout_rate: 0.0,
    };

    let network = RealNeuralNetwork::new(model_config, device.clone())?;

    let input_data = vec![1.0f32; 21];
    let input = Tensor::from_vec(input_data, &[1, 21], &device)?;

    // Measure latency for first 10 iterations
    let num_iterations = 10;
    let mut latencies = Vec::new();

    for i in 0..num_iterations {
        let start = Instant::now();
        let _ = network.forward(&input).await?;
        let latency = start.elapsed().as_micros();
        latencies.push(latency);

        println!("Iteration {}: {}μs", i + 1, latency);
    }

    // Later iterations should stabilize
    let first_half_avg = latencies[0..5].iter().sum::<u128>() / 5;
    let second_half_avg = latencies[5..10].iter().sum::<u128>() / 5;

    println!("First half average: {}μs", first_half_avg);
    println!("Second half average: {}μs", second_half_avg);

    // Verify all latencies are positive
    assert!(latencies.iter().all(|&l| l > 0));

    Ok(())
}

// ==================== THROUGHPUT BENCHMARKS ====================

/// Test: Maximum throughput measurement
#[tokio::test]
#[ignore = "Slow test - run explicitly for benchmarking"]
async fn test_maximum_throughput() -> Result<(), Box<dyn std::error::Error>> {
    let safety_manager = Arc::new(MLSafetyManager::new(MLSafetyConfig::default()));
    let mut config = RealInferenceConfig::default();
    config.device_preference = "cpu".to_string();
    config.enable_caching = false;

    let engine = RealMLInferenceEngine::new(config, safety_manager);

    let model_config = ModelConfig {
        input_dim: 21,
        hidden_dims: vec![32],
        output_dim: 1,
        activation: "relu".to_string(),
        batch_norm: false,
        dropout_rate: 0.0,
    };

    engine
        .load_model("throughput_test".to_string(), model_config)
        .await?;

    // Warmup
    let features = create_mock_features();
    let _ = engine.predict("throughput_test", &features).await;

    // Measure max throughput for 5 seconds
    let duration = Duration::from_secs(5);
    let start = Instant::now();
    let mut num_inferences = 0;

    while start.elapsed() < duration {
        let features = create_mock_features();
        let _ = engine.predict("throughput_test", &features).await;
        num_inferences += 1;
    }

    let elapsed_secs = start.elapsed().as_secs_f64();
    let throughput = num_inferences as f64 / elapsed_secs;

    println!(
        "Maximum throughput: {:.2} inferences/sec ({} inferences in {:.3}s)",
        throughput, num_inferences, elapsed_secs
    );

    // Target: >1000 inferences/sec (CPU might be slower)
    assert!(throughput > 0.0);

    Ok(())
}

/// Test: Sustained load - no degradation
#[tokio::test]
#[ignore = "Slow test - run explicitly for stress testing"]
async fn test_sustained_load_no_degradation() -> Result<(), Box<dyn std::error::Error>> {
    let safety_manager = Arc::new(MLSafetyManager::new(MLSafetyConfig::default()));
    let mut config = RealInferenceConfig::default();
    config.device_preference = "cpu".to_string();

    let engine = RealMLInferenceEngine::new(config, safety_manager);

    let model_config = ModelConfig {
        input_dim: 21,
        hidden_dims: vec![32],
        output_dim: 1,
        activation: "relu".to_string(),
        batch_norm: false,
        dropout_rate: 0.0,
    };

    engine
        .load_model("sustained_test".to_string(), model_config)
        .await?;

    // Run for 1 minute, measure latency every 10 seconds
    let total_duration = Duration::from_secs(60);
    let measure_interval = Duration::from_secs(10);
    let start = Instant::now();

    while start.elapsed() < total_duration {
        let interval_start = Instant::now();
        let mut interval_latencies = Vec::new();

        // Measure for one interval
        while interval_start.elapsed() < measure_interval {
            let features = create_mock_features();
            let measure_start = Instant::now();
            let _ = engine.predict("sustained_test", &features).await?;
            interval_latencies.push(measure_start.elapsed().as_micros());
        }

        if !interval_latencies.is_empty() {
            let avg_latency =
                interval_latencies.iter().sum::<u128>() / interval_latencies.len() as u128;
            let elapsed_secs = start.elapsed().as_secs();
            println!(
                "At {}s: avg latency {}μs ({} samples)",
                elapsed_secs,
                avg_latency,
                interval_latencies.len()
            );
        }
    }

    println!("Sustained load test completed");

    Ok(())
}