foxhunt/ml/tests/memory_optimization_tests.rs

//! Comprehensive Memory Optimization Tests for 4GB GPU
//!
//! Tests quantization, mixed precision, and memory efficiency features
//! to ensure training fits within RTX 3050 Ti 4GB VRAM constraints.

use candle_core::{DType, Device, Tensor};
use ml::memory_optimization::{
    MemoryOptimizationConfig, MemoryStats, PrecisionConverter, PrecisionType, QuantizationConfig,
    QuantizationType, Quantizer,
};

/// Helper to create test device (CUDA if available, CPU fallback)
fn test_device() -> Device {
    Device::cuda_if_available(0).unwrap_or(Device::Cpu)
}

/// Helper to create test tensor
fn create_test_tensor(device: &Device, shape: &[usize]) -> Tensor {
    Tensor::randn(0.0f32, 1.0f32, shape, device).unwrap()
}

#[test]
fn test_int8_quantization_basic() {
    let device = test_device();
    println!("Running INT8 quantization test on {:?}", device);

    // Create test tensor
    let tensor = create_test_tensor(&device, &[256, 256]);
    let original_size = tensor.dims().iter().product::<usize>() * 4; // 4 bytes per f32

    println!(
        "Original tensor: {:?}, size: {} bytes",
        tensor.dims(),
        original_size
    );

    // Configure INT8 quantization
    let config = QuantizationConfig {
        quant_type: QuantizationType::Int8,
        symmetric: true,
        per_channel: true,
        calibration_samples: Some(1000),
    };

    let mut quantizer = Quantizer::new(config, device.clone());

    // Quantize tensor
    let quantized = quantizer
        .quantize_tensor(&tensor, "test_layer")
        .expect("Quantization failed");

    println!(
        "Quantized type: {:?}, scale: {}, zero_point: {}",
        quantized.quant_type, quantized.scale, quantized.zero_point
    );

    // Verify quantization type
    assert_eq!(quantized.quant_type, QuantizationType::Int8);

    // Check memory savings
    let quantized_size = quantized.memory_bytes();
    let savings_percent = (1.0 - (quantized_size as f64 / original_size as f64)) * 100.0;

    println!(
        "Original: {} bytes, Quantized: {} bytes, Savings: {:.1}%",
        original_size, quantized_size, savings_percent
    );

    // INT8 should achieve ~75% memory reduction
    assert!(
        savings_percent >= 70.0,
        "Expected at least 70% memory savings"
    );

    // Dequantize and check accuracy
    let dequantized = quantizer
        .dequantize_tensor(&quantized)
        .expect("Dequantization failed");

    assert_eq!(dequantized.dims(), tensor.dims());
    println!("✓ INT8 quantization test passed");
}

#[test]
fn test_int4_quantization() {
    let device = test_device();
    println!("Running INT4 quantization test on {:?}", device);

    let tensor = create_test_tensor(&device, &[512, 512]);
    let original_size = tensor.dims().iter().product::<usize>() * 4;

    let config = QuantizationConfig {
        quant_type: QuantizationType::Int4,
        symmetric: true,
        per_channel: false,
        calibration_samples: None,
    };

    let mut quantizer = Quantizer::new(config, device.clone());
    let quantized = quantizer
        .quantize_tensor(&tensor, "test_layer_int4")
        .expect("INT4 quantization failed");

    assert_eq!(quantized.quant_type, QuantizationType::Int4);

    let quantized_size = quantized.memory_bytes();
    let savings_percent = (1.0 - (quantized_size as f64 / original_size as f64)) * 100.0;

    println!(
        "INT4 - Original: {} bytes, Quantized: {} bytes, Savings: {:.1}%",
        original_size, quantized_size, savings_percent
    );

    // INT4 should achieve ~87.5% memory reduction
    assert!(
        savings_percent >= 85.0,
        "Expected at least 85% memory savings"
    );
    println!("✓ INT4 quantization test passed");
}

#[test]
fn test_asymmetric_quantization() {
    let device = test_device();
    println!("Running asymmetric quantization test on {:?}", device);

    let tensor = create_test_tensor(&device, &[128, 128]);

    let config = QuantizationConfig {
        quant_type: QuantizationType::Int8,
        symmetric: false, // Asymmetric
        per_channel: true,
        calibration_samples: Some(500),
    };

    let mut quantizer = Quantizer::new(config, device.clone());
    let quantized = quantizer
        .quantize_tensor(&tensor, "asymmetric_layer")
        .expect("Asymmetric quantization failed");

    // Asymmetric quantization should use non-zero zero_point
    println!(
        "Asymmetric quantization - scale: {}, zero_point: {}",
        quantized.scale, quantized.zero_point
    );

    assert_eq!(quantized.quant_type, QuantizationType::Int8);
    println!("✓ Asymmetric quantization test passed");
}

#[test]
fn test_float16_precision_conversion() {
    let device = test_device();
    println!("Running FP16 precision test on {:?}", device);

    let tensor = create_test_tensor(&device, &[256, 256]);
    let original_size = tensor.dims().iter().product::<usize>() * 4; // F32

    let mut converter = PrecisionConverter::new(PrecisionType::Float16, device.clone());

    let converted = converter
        .to_float16(&tensor)
        .expect("FP16 conversion failed");

    assert_eq!(converted.dtype(), DType::F16);

    let converted_size = converted.dims().iter().product::<usize>() * 2; // F16 = 2 bytes
    let savings_percent = (1.0 - (converted_size as f64 / original_size as f64)) * 100.0;

    println!(
        "FP16 - Original: {} bytes (F32), Converted: {} bytes (F16), Savings: {:.1}%",
        original_size, converted_size, savings_percent
    );

    // FP16 should achieve 50% memory reduction
    assert!(savings_percent >= 49.0 && savings_percent <= 51.0);

    // Check statistics
    let stats = converter.get_stats();
    println!(
        "Conversion stats: {} conversions, {:.2} MB saved",
        stats.conversions, stats.memory_saved_mb
    );

    assert_eq!(stats.conversions, 1);
    assert!(stats.memory_saved_mb > 0.0);
    println!("✓ FP16 precision conversion test passed");
}

#[test]
fn test_bfloat16_precision_conversion() {
    let device = test_device();
    println!("Running BF16 precision test on {:?}", device);

    let tensor = create_test_tensor(&device, &[512, 512]);

    let mut converter = PrecisionConverter::new(PrecisionType::BFloat16, device.clone());

    let converted = converter
        .to_bfloat16(&tensor)
        .expect("BF16 conversion failed");

    assert_eq!(converted.dtype(), DType::BF16);

    let original_size = tensor.dims().iter().product::<usize>() * 4;
    let converted_size = converted.dims().iter().product::<usize>() * 2;
    let savings_percent = (1.0 - (converted_size as f64 / original_size as f64)) * 100.0;

    println!(
        "BF16 - Original: {} bytes (F32), Converted: {} bytes (BF16), Savings: {:.1}%",
        original_size, converted_size, savings_percent
    );

    assert!(savings_percent >= 49.0 && savings_percent <= 51.0);
    println!("✓ BF16 precision conversion test passed");
}

#[test]
fn test_mixed_precision_roundtrip() {
    let device = test_device();
    println!("Running mixed precision roundtrip test on {:?}", device);

    let original = create_test_tensor(&device, &[128, 128]);

    let mut converter = PrecisionConverter::new(PrecisionType::Float16, device.clone());

    // Convert F32 -> F16 -> F32
    let fp16 = converter.to_float16(&original).expect("F32->F16 failed");
    let restored = converter.to_float32(&fp16).expect("F16->F32 failed");

    assert_eq!(restored.dtype(), DType::F32);
    assert_eq!(restored.dims(), original.dims());

    // Validate accuracy
    let accuracy =
        ml::memory_optimization::precision::validate_precision_accuracy(&original, &restored)
            .expect("Accuracy validation failed");

    println!(
        "Accuracy metrics: MAE={:.6}, RMSE={:.6}, Relative Error={:.6}%",
        accuracy.mae,
        accuracy.rmse,
        accuracy.mean_relative_error * 100.0
    );

    // FP16 should maintain reasonable accuracy (<5% error)
    assert!(
        accuracy.is_acceptable(5.0),
        "Relative error too high: {:.2}%",
        accuracy.mean_relative_error * 100.0
    );

    println!("✓ Mixed precision roundtrip test passed");
}

#[test]
fn test_quantization_accuracy_preservation() {
    let device = test_device();
    println!("Running quantization accuracy test on {:?}", device);

    let original = create_test_tensor(&device, &[256, 256]);

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

    let mut quantizer = Quantizer::new(config, device.clone());

    // Quantize and dequantize
    let quantized = quantizer
        .quantize_tensor(&original, "accuracy_test")
        .expect("Quantization failed");

    let restored = quantizer
        .dequantize_tensor(&quantized)
        .expect("Dequantization failed");

    // Validate accuracy
    let accuracy =
        ml::memory_optimization::precision::validate_precision_accuracy(&original, &restored)
            .expect("Accuracy validation failed");

    println!(
        "Quantization accuracy: MAE={:.6}, RMSE={:.6}, Max Error={:.6}",
        accuracy.mae, accuracy.rmse, accuracy.max_absolute_error
    );

    // INT8 quantization should maintain reasonable accuracy
    assert!(accuracy.rmse < 0.1, "RMSE too high: {:.6}", accuracy.rmse);

    println!("✓ Quantization accuracy preservation test passed");
}

#[test]
fn test_memory_optimization_config() {
    println!("Testing memory optimization configuration");

    let config = MemoryOptimizationConfig::default();

    assert!(config.lazy_loading);
    assert_eq!(config.precision, PrecisionType::Float32);
    assert_eq!(config.quantization, QuantizationType::None);
    assert!(config.tensor_caching);

    // Custom config for 4GB GPU
    let custom_config = MemoryOptimizationConfig {
        lazy_loading: true,
        precision: PrecisionType::Float16,
        quantization: QuantizationType::Int8,
        max_memory_mb: Some(3500.0), // Leave 500MB headroom
        gradient_checkpointing: true,
        tensor_caching: false, // Reduce cache memory
    };

    assert_eq!(custom_config.precision, PrecisionType::Float16);
    assert_eq!(custom_config.quantization, QuantizationType::Int8);
    assert_eq!(custom_config.max_memory_mb, Some(3500.0));

    println!("✓ Memory optimization config test passed");
}

#[test]
fn test_memory_stats_tracking() {
    println!("Testing memory statistics tracking");

    let mut stats = MemoryStats::new();

    assert_eq!(stats.current_mb, 0.0);
    assert_eq!(stats.peak_mb, 0.0);

    // Simulate memory usage
    stats.update_peak(100.0);
    assert_eq!(stats.current_mb, 100.0);
    assert_eq!(stats.peak_mb, 100.0);

    stats.update_peak(150.0);
    assert_eq!(stats.current_mb, 150.0);
    assert_eq!(stats.peak_mb, 150.0);

    stats.update_peak(120.0); // Peak should not decrease
    assert_eq!(stats.current_mb, 120.0);
    assert_eq!(stats.peak_mb, 150.0);

    // Add component breakdown
    stats.add_component("model_weights", 50.0);
    stats.add_component("activations", 30.0);
    stats.add_component("optimizer_state", 20.0);

    assert_eq!(stats.breakdown.len(), 3);
    assert_eq!(stats.breakdown.get("model_weights"), Some(&50.0));

    println!("✓ Memory stats tracking test passed");
}

#[test]
fn test_multi_tensor_quantization() {
    let device = test_device();
    println!("Running multi-tensor quantization test on {:?}", device);

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

    let mut quantizer = Quantizer::new(config, device.clone());

    // Quantize multiple tensors (simulating model layers)
    let tensors = vec![
        create_test_tensor(&device, &[256, 256]),
        create_test_tensor(&device, &[512, 512]),
        create_test_tensor(&device, &[1024, 256]),
        create_test_tensor(&device, &[256, 128]),
    ];

    let layer_names = vec!["layer1", "layer2", "layer3", "layer4"];

    for (tensor, name) in tensors.iter().zip(layer_names.iter()) {
        let quantized = quantizer
            .quantize_tensor(tensor, name)
            .expect("Multi-tensor quantization failed");

        println!("Quantized {}: {} bytes", name, quantized.memory_bytes());
    }

    // Check total memory savings
    let savings_mb = quantizer.memory_savings_mb();
    println!("Total memory savings: {:.2} MB", savings_mb);

    assert!(savings_mb > 0.0, "No memory savings recorded");
    println!("✓ Multi-tensor quantization test passed");
}

#[test]
fn test_precision_converter_stats() {
    let device = test_device();
    println!("Testing precision converter statistics on {:?}", device);

    let mut converter = PrecisionConverter::new(PrecisionType::Float16, device.clone());

    // Convert multiple tensors
    for i in 0..5 {
        let tensor = create_test_tensor(&device, &[128, 128]);
        let _converted = converter.to_float16(&tensor).expect("Conversion failed");
        println!("Converted tensor {}/5", i + 1);
    }

    let stats = converter.get_stats();

    assert_eq!(stats.conversions, 5);
    assert!(stats.memory_saved_mb > 0.0);
    assert_eq!(stats.target_precision, PrecisionType::Float16);

    println!(
        "Stats: {} conversions, {:.2} MB saved",
        stats.conversions, stats.memory_saved_mb
    );

    // Reset and verify
    converter.reset_stats();
    let new_stats = converter.get_stats();
    assert_eq!(new_stats.conversions, 0);
    assert_eq!(new_stats.memory_saved_mb, 0.0);

    println!("✓ Precision converter stats test passed");
}

#[test]
fn test_4gb_gpu_memory_compatibility() {
    let device = test_device();
    println!("Testing 4GB GPU memory compatibility on {:?}", device);

    // Simulate MAMBA-2 model sizes with memory optimization
    let model_configs = vec![
        (
            "baseline_f32",
            4,
            QuantizationType::None,
            PrecisionType::Float32,
        ),
        (
            "int8_f32",
            4,
            QuantizationType::Int8,
            PrecisionType::Float32,
        ),
        (
            "none_f16",
            4,
            QuantizationType::None,
            PrecisionType::Float16,
        ),
        (
            "int8_f16",
            4,
            QuantizationType::Int8,
            PrecisionType::Float16,
        ),
    ];

    for (name, size_multiplier, quant_type, precision) in model_configs {
        // Estimate memory usage for different configs
        let base_size_mb = 500.0; // MAMBA-2 base size
        let model_size = base_size_mb * size_multiplier as f64;

        let memory_multiplier = precision.memory_multiplier();
        let quant_savings = match quant_type {
            QuantizationType::None => 1.0,
            QuantizationType::Int8 => 0.25,
            QuantizationType::Int4 => 0.125,
            QuantizationType::Dynamic => 0.25,
        };

        let final_size = model_size * memory_multiplier * quant_savings;
        let fits_4gb = final_size <= 3500.0; // Leave 500MB headroom

        println!(
            "Config '{}': {:.1} MB (quant={:?}, precision={:?}) - {}",
            name,
            final_size,
            quant_type,
            precision,
            if fits_4gb {
                "✓ FITS"
            } else {
                "✗ TOO LARGE"
            }
        );
    }

    println!("✓ 4GB GPU memory compatibility test passed");
}

#[test]
fn test_gradient_checkpointing_simulation() {
    println!("Testing gradient checkpointing simulation");

    let config = MemoryOptimizationConfig {
        lazy_loading: true,
        precision: PrecisionType::Float32,
        quantization: QuantizationType::None,
        max_memory_mb: Some(3500.0),
        gradient_checkpointing: true,
        tensor_caching: false,
    };

    assert!(config.gradient_checkpointing);

    // Gradient checkpointing typically reduces activation memory by ~2-3x
    // at the cost of ~33% more compute time
    let activation_memory_mb = 1000.0;
    let with_checkpointing = activation_memory_mb / 2.5;
    let savings = activation_memory_mb - with_checkpointing;

    println!(
        "Gradient checkpointing: {:.1} MB -> {:.1} MB (saves {:.1} MB)",
        activation_memory_mb, with_checkpointing, savings
    );

    assert!(savings > 0.0);
    println!("✓ Gradient checkpointing simulation test passed");
}

#[test]
fn test_no_quantization_passthrough() {
    let device = test_device();
    println!("Testing no-quantization passthrough on {:?}", device);

    let tensor = create_test_tensor(&device, &[128, 128]);
    let original_size = tensor.dims().iter().product::<usize>() * 4;

    let config = QuantizationConfig {
        quant_type: QuantizationType::None,
        symmetric: true,
        per_channel: false,
        calibration_samples: None,
    };

    let mut quantizer = Quantizer::new(config, device.clone());

    let result = quantizer
        .quantize_tensor(&tensor, "passthrough_test")
        .expect("Passthrough failed");

    assert_eq!(result.quant_type, QuantizationType::None);
    assert_eq!(result.memory_bytes(), original_size);

    println!("✓ No-quantization passthrough test passed");
}

#[test]
fn test_precision_type_properties() {
    println!("Testing precision type properties");

    assert_eq!(PrecisionType::Float32.bytes_per_element(), 4);
    assert_eq!(PrecisionType::Float16.bytes_per_element(), 2);
    assert_eq!(PrecisionType::BFloat16.bytes_per_element(), 2);

    assert_eq!(PrecisionType::Float32.memory_multiplier(), 1.0);
    assert_eq!(PrecisionType::Float16.memory_multiplier(), 0.5);
    assert_eq!(PrecisionType::BFloat16.memory_multiplier(), 0.5);

    assert_eq!(PrecisionType::Float32.to_dtype(), DType::F32);
    assert_eq!(PrecisionType::Float16.to_dtype(), DType::F16);
    assert_eq!(PrecisionType::BFloat16.to_dtype(), DType::BF16);

    println!("✓ Precision type properties test passed");
}

#[test]
fn test_memory_optimization_full_pipeline() {
    let device = test_device();
    println!(
        "Running full memory optimization pipeline test on {:?}",
        device
    );

    let mut stats = MemoryStats::new();

    // Step 1: Create baseline model (F32)
    let model_tensor = create_test_tensor(&device, &[512, 512]);
    let baseline_size = (model_tensor.dims().iter().product::<usize>() * 4) as f64 / 1_048_576.0;
    stats.add_component("baseline_model", baseline_size);
    stats.update_peak(baseline_size);

    println!("Step 1: Baseline model (F32): {:.2} MB", baseline_size);

    // Step 2: Apply FP16 precision
    let mut precision_converter = PrecisionConverter::new(PrecisionType::Float16, device.clone());
    let fp16_tensor = precision_converter
        .to_float16(&model_tensor)
        .expect("FP16 conversion failed");

    let fp16_size = (fp16_tensor.dims().iter().product::<usize>() * 2) as f64 / 1_048_576.0;
    let precision_savings = baseline_size - fp16_size;
    stats.add_component("fp16_model", fp16_size);
    stats.savings_mb += precision_savings;

    println!(
        "Step 2: FP16 model: {:.2} MB (saved {:.2} MB)",
        fp16_size, precision_savings
    );

    // Step 3: Apply INT8 quantization
    let fp32_for_quant = precision_converter
        .to_float32(&fp16_tensor)
        .expect("F32 conversion failed");

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

    let mut quantizer = Quantizer::new(quant_config, device.clone());
    let quantized = quantizer
        .quantize_tensor(&fp32_for_quant, "optimized_model")
        .expect("Quantization failed");

    let quantized_size = quantized.memory_bytes() as f64 / 1_048_576.0;
    let quant_savings = fp16_size - quantized_size;
    stats.add_component("int8_fp16_model", quantized_size);
    stats.savings_mb += quant_savings;

    println!(
        "Step 3: INT8+FP16 model: {:.2} MB (saved {:.2} MB)",
        quantized_size, quant_savings
    );

    // Final results
    let total_savings = baseline_size - quantized_size;
    let savings_percent = (total_savings / baseline_size) * 100.0;

    println!("\n=== Memory Optimization Summary ===");
    println!("Baseline (F32):      {:.2} MB", baseline_size);
    println!("Optimized (INT8+FP16): {:.2} MB", quantized_size);
    println!(
        "Total Savings:       {:.2} MB ({:.1}%)",
        total_savings, savings_percent
    );
    println!(
        "Fits in 4GB GPU:     {}",
        if quantized_size < 3500.0 {
            "✓ YES"
        } else {
            "✗ NO"
        }
    );

    // Verify significant savings
    assert!(
        savings_percent >= 85.0,
        "Expected at least 85% memory savings"
    );

    println!("✓ Full memory optimization pipeline test passed");
}