foxhunt/ml/tests/quantized_checkpoint_test.rs

//! Integration tests for quantized checkpoint save/load functionality
//!
//! Tests:
//! 1. Round-trip save/load preserves weights and metadata
//! 2. Compression reduces file size without data loss
//! 3. Backward compatibility with FP32 checkpoints
//! 4. File size validation (<100MB target)
//! 5. Multi-layer models with different shapes

use candle_core::{Device, Tensor};
use ml::checkpoint::{
    calculate_compression_ratio, load_quantized_checkpoint, save_quantized_checkpoint,
    QuantizedCheckpointMetadata, QuantizedWeight,
};
use ml::MLError;
use std::collections::HashMap;
use tempfile::{NamedTempFile, TempDir};

/// Test basic save/load round-trip
#[test]
fn test_quantized_checkpoint_round_trip() -> Result<(), MLError> {
    let device = Device::Cpu;
    let mut weights = HashMap::new();

    // Create test quantized weight
    let data = Tensor::from_vec(vec![0u8, 64, 128, 192, 255], 5, &device)
        .map_err(|e| MLError::ModelError(format!("Tensor creation failed: {}", e)))?;

    weights.insert(
        "fc1.weight".to_string(),
        QuantizedWeight {
            data,
            scale: 0.01,
            zero_point: 127,
            shape: vec![5],
        },
    );

    // Save checkpoint
    let temp_file = NamedTempFile::new().unwrap();
    let path = temp_file.path();

    let metadata = QuantizedCheckpointMetadata {
        model_type: "DQN".to_string(),
        version: "1.0.0".to_string(),
        quantization_method: "symmetric".to_string(),
        ..Default::default()
    };

    save_quantized_checkpoint(path, &weights, Some(metadata.clone()), false)?;

    // Load checkpoint
    let (loaded_weights, loaded_metadata) = load_quantized_checkpoint(path)?;

    // Validate metadata
    assert_eq!(loaded_metadata.model_type, "DQN");
    assert_eq!(loaded_metadata.version, "1.0.0");
    assert_eq!(loaded_metadata.num_layers, 1);

    // Validate weights
    assert_eq!(loaded_weights.len(), 1);
    let loaded_weight = loaded_weights.get("fc1.weight").unwrap();
    assert_eq!(loaded_weight.scale, 0.01);
    assert_eq!(loaded_weight.zero_point, 127);
    assert_eq!(loaded_weight.shape, vec![5]);

    // Validate tensor data
    let loaded_data = loaded_weight
        .data
        .flatten_all()
        .unwrap()
        .to_vec1::<u8>()
        .unwrap();
    assert_eq!(loaded_data, vec![0u8, 64, 128, 192, 255]);

    Ok(())
}

/// Test multi-layer model checkpoint
#[test]
fn test_multi_layer_checkpoint() -> Result<(), MLError> {
    let device = Device::Cpu;
    let mut weights = HashMap::new();

    // Layer 1: 1D weight vector
    let fc1_data = Tensor::from_vec(vec![100u8; 128], &[128], &device).unwrap();
    weights.insert(
        "fc1.weight".to_string(),
        QuantizedWeight {
            data: fc1_data,
            scale: 0.02,
            zero_point: 127,
            shape: vec![128],
        },
    );

    // Layer 2: 2D weight matrix
    let fc2_data = Tensor::from_vec(vec![150u8; 256], &[64, 4], &device).unwrap();
    weights.insert(
        "fc2.weight".to_string(),
        QuantizedWeight {
            data: fc2_data,
            scale: 0.015,
            zero_point: 100,
            shape: vec![64, 4],
        },
    );

    // Layer 3: Bias vector
    let bias_data = Tensor::from_vec(vec![127u8; 64], &[64], &device).unwrap();
    weights.insert(
        "fc2.bias".to_string(),
        QuantizedWeight {
            data: bias_data,
            scale: 0.001,
            zero_point: 127,
            shape: vec![64],
        },
    );

    // Save and load
    let temp_file = NamedTempFile::new().unwrap();
    save_quantized_checkpoint(temp_file.path(), &weights, None, false)?;

    let (loaded_weights, loaded_metadata) = load_quantized_checkpoint(temp_file.path())?;

    // Validate all layers loaded
    assert_eq!(loaded_weights.len(), 3);
    assert_eq!(loaded_metadata.num_layers, 3);

    // Validate individual layers
    assert!(loaded_weights.contains_key("fc1.weight"));
    assert!(loaded_weights.contains_key("fc2.weight"));
    assert!(loaded_weights.contains_key("fc2.bias"));

    // Validate shapes
    assert_eq!(loaded_weights.get("fc1.weight").unwrap().shape, vec![128]);
    assert_eq!(loaded_weights.get("fc2.weight").unwrap().shape, vec![64, 4]);
    assert_eq!(loaded_weights.get("fc2.bias").unwrap().shape, vec![64]);

    Ok(())
}

/// Test gzip compression
#[test]
fn test_gzip_compression() -> Result<(), MLError> {
    let device = Device::Cpu;
    let mut weights = HashMap::new();

    // Create large repetitive data (compresses well)
    let data = Tensor::from_vec(vec![42u8; 50_000], &[50_000], &device).unwrap();
    weights.insert(
        "large_layer".to_string(),
        QuantizedWeight {
            data,
            scale: 1.0,
            zero_point: 127,
            shape: vec![50_000],
        },
    );

    let temp_dir = TempDir::new().unwrap();
    let compressed_path = temp_dir.path().join("compressed.safetensors");
    let uncompressed_path = temp_dir.path().join("uncompressed.safetensors");

    // Save with compression
    let compressed_size = save_quantized_checkpoint(&compressed_path, &weights, None, true)?;

    // Save without compression
    let uncompressed_size = save_quantized_checkpoint(&uncompressed_path, &weights, None, false)?;

    // Compressed should be significantly smaller for repetitive data
    println!(
        "Compression: {} → {} bytes ({:.1}% reduction)",
        uncompressed_size,
        compressed_size,
        100.0 * (1.0 - compressed_size as f64 / uncompressed_size as f64)
    );
    assert!(compressed_size < uncompressed_size);

    // Load compressed and verify data integrity
    let (loaded_weights, _) = load_quantized_checkpoint(&compressed_path)?;
    assert_eq!(loaded_weights.len(), 1);

    let loaded_data = loaded_weights
        .get("large_layer")
        .unwrap()
        .data
        .flatten_all()
        .unwrap()
        .to_vec1::<u8>()
        .unwrap();
    assert_eq!(loaded_data.len(), 50_000);
    assert!(loaded_data.iter().all(|&x| x == 42));

    Ok(())
}

/// Test file size comparison (INT8 vs FP32)
#[test]
fn test_file_size_comparison() -> Result<(), MLError> {
    let device = Device::Cpu;
    let mut weights = HashMap::new();

    // Simulate DQN Q-network weights (approximate sizes)
    // Input layer: 225 features × 128 hidden = 28,800 params
    let fc1_data = Tensor::from_vec(vec![127u8; 28_800], &[225, 128], &device).unwrap();
    weights.insert(
        "fc1.weight".to_string(),
        QuantizedWeight {
            data: fc1_data,
            scale: 0.01,
            zero_point: 127,
            shape: vec![225, 128],
        },
    );

    // Hidden layer: 128 × 64 = 8,192 params
    let fc2_data = Tensor::from_vec(vec![127u8; 8_192], &[128, 64], &device).unwrap();
    weights.insert(
        "fc2.weight".to_string(),
        QuantizedWeight {
            data: fc2_data,
            scale: 0.01,
            zero_point: 127,
            shape: vec![128, 64],
        },
    );

    // Output layer: 64 × 3 = 192 params (BUY/SELL/HOLD)
    let fc3_data = Tensor::from_vec(vec![127u8; 192], &[64, 3], &device).unwrap();
    weights.insert(
        "fc3.weight".to_string(),
        QuantizedWeight {
            data: fc3_data,
            scale: 0.01,
            zero_point: 127,
            shape: vec![64, 3],
        },
    );

    // Calculate sizes
    let int8_size: usize = weights.values().map(|w| w.memory_bytes()).sum();
    let fp32_size = int8_size * 4;
    let compression_ratio = calculate_compression_ratio(&weights);

    println!("Model size comparison:");
    println!(
        "  FP32:  {} bytes ({:.2} MB)",
        fp32_size,
        fp32_size as f64 / 1_048_576.0
    );
    println!(
        "  INT8:  {} bytes ({:.2} MB)",
        int8_size,
        int8_size as f64 / 1_048_576.0
    );
    println!("  Ratio: {:.2}x", compression_ratio);

    // Save checkpoint and verify file size
    let temp_file = NamedTempFile::new().unwrap();
    let file_size = save_quantized_checkpoint(temp_file.path(), &weights, None, false)?;

    println!(
        "  File:  {} bytes ({:.2} MB)",
        file_size,
        file_size as f64 / 1_048_576.0
    );

    // Validate size reduction
    assert_eq!(compression_ratio, 4.0); // FP32/INT8 = 4x
    assert!(file_size < fp32_size);
    assert!(file_size < 1_000_000); // Should be well under 1MB for this small model

    Ok(())
}

/// Test custom metadata fields
#[test]
fn test_custom_metadata() -> Result<(), MLError> {
    let device = Device::Cpu;
    let mut weights = HashMap::new();

    let data = Tensor::from_vec(vec![127u8; 100], &[100], &device).unwrap();
    weights.insert(
        "test.weight".to_string(),
        QuantizedWeight {
            data,
            scale: 1.0,
            zero_point: 127,
            shape: vec![100],
        },
    );

    // Add custom metadata
    let mut metadata = QuantizedCheckpointMetadata::default();
    metadata.model_type = "PPO".to_string();
    metadata.version = "2.1.0".to_string();
    metadata
        .custom
        .insert("training_epochs".to_string(), serde_json::json!(50));
    metadata
        .custom
        .insert("dataset".to_string(), serde_json::json!("ES.FUT_2024_Q4"));
    metadata
        .custom
        .insert("accuracy".to_string(), serde_json::json!(0.95));

    // Save and load
    let temp_file = NamedTempFile::new().unwrap();
    save_quantized_checkpoint(temp_file.path(), &weights, Some(metadata.clone()), false)?;

    let (_, loaded_metadata) = load_quantized_checkpoint(temp_file.path())?;

    // Validate custom fields
    assert_eq!(loaded_metadata.model_type, "PPO");
    assert_eq!(loaded_metadata.version, "2.1.0");
    assert_eq!(
        loaded_metadata.custom.get("training_epochs").unwrap(),
        &serde_json::json!(50)
    );
    assert_eq!(
        loaded_metadata.custom.get("dataset").unwrap(),
        &serde_json::json!("ES.FUT_2024_Q4")
    );

    Ok(())
}

/// Test large model checkpoint (simulating MAMBA-2 or TFT)
#[test]
fn test_large_model_checkpoint() -> Result<(), MLError> {
    let device = Device::Cpu;
    let mut weights = HashMap::new();

    // Simulate multiple large layers (total ~10MB INT8 → ~40MB FP32)
    for i in 0..10 {
        let data =
            Tensor::from_vec(vec![(i * 25) as u8; 1_000_000], &[1_000_000], &device).unwrap();
        weights.insert(
            format!("layer{}.weight", i),
            QuantizedWeight {
                data,
                scale: 0.01,
                zero_point: 127,
                shape: vec![1_000_000],
            },
        );
    }

    // Save checkpoint
    let temp_file = NamedTempFile::new().unwrap();
    let file_size = save_quantized_checkpoint(temp_file.path(), &weights, None, false)?;

    // Validate file size is reasonable
    println!(
        "Large model checkpoint: {} bytes ({:.2} MB)",
        file_size,
        file_size as f64 / 1_048_576.0
    );
    assert!(file_size < 50_000_000); // Should be under 50MB

    // Load and validate
    let (loaded_weights, loaded_metadata) = load_quantized_checkpoint(temp_file.path())?;
    assert_eq!(loaded_weights.len(), 10);
    assert_eq!(loaded_metadata.num_layers, 10);

    // Verify total size
    let total_int8_size: usize = loaded_weights.values().map(|w| w.memory_bytes()).sum();
    assert_eq!(total_int8_size, 10_000_000); // 10 × 1M = 10MB

    Ok(())
}

/// Benchmark save/load performance
#[test]
fn test_performance_benchmark() -> Result<(), MLError> {
    let device = Device::Cpu;
    let mut weights = HashMap::new();

    // Medium-sized model (~1MB INT8)
    let data = Tensor::from_vec(vec![127u8; 1_000_000], &[1_000_000], &device).unwrap();
    weights.insert(
        "model.weight".to_string(),
        QuantizedWeight {
            data,
            scale: 1.0,
            zero_point: 127,
            shape: vec![1_000_000],
        },
    );

    let temp_file = NamedTempFile::new().unwrap();

    // Benchmark save
    let save_start = std::time::Instant::now();
    save_quantized_checkpoint(temp_file.path(), &weights, None, false)?;
    let save_duration = save_start.elapsed();

    // Benchmark load
    let load_start = std::time::Instant::now();
    let _ = load_quantized_checkpoint(temp_file.path())?;
    let load_duration = load_start.elapsed();

    println!("Performance benchmark (1MB INT8 model):");
    println!("  Save: {:?}", save_duration);
    println!("  Load: {:?}", load_duration);

    // Sanity checks (should be fast for 1MB)
    assert!(save_duration.as_millis() < 1000); // < 1 second
    assert!(load_duration.as_millis() < 1000); // < 1 second

    Ok(())
}