foxhunt/ml/tests/qat_test.rs

//! Comprehensive Unit Tests for Quantization-Aware Training (QAT)
//!
//! Tests all aspects of QAT implementation:
//! 1. Fake quantization forward pass (quantize→dequantize round-trip)
//! 2. Gradient flow through fake quantization (Straight-Through Estimator)
//! 3. Observer statistics tracking (min/max with EMA)
//! 4. QAT calibration phase workflow
//! 5. QAT→INT8 conversion for deployment
//! 6. QAT vs PTQ accuracy comparison (1-2% improvement expected)

use candle_core::{DType, Device, Tensor};
use ml::memory_optimization::{
    compare_qat_vs_ptq_accuracy, FakeQuantize, QATConfig, QuantizationConfig, QuantizationObserver,
    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 with known range
fn create_test_tensor(device: &Device, shape: &[usize]) -> Tensor {
    Tensor::randn(0.0f32, 1.0f32, shape, device).unwrap()
}

/// Helper to create calibration data (multiple batches)
fn create_calibration_data(
    device: &Device,
    num_batches: usize,
    batch_shape: &[usize],
) -> Vec<Tensor> {
    (0..num_batches)
        .map(|_| create_test_tensor(device, batch_shape))
        .collect()
}

// ============================================================================
// Test 1: Fake Quantize Forward Pass (Quantize→Dequantize Round-Trip)
// ============================================================================

#[test]
fn test_fake_quantize_forward() {
    println!("\n=== Test 1: Fake Quantize Forward Pass ===");
    let device = test_device();
    println!("Device: {:?}", device);

    // Create test tensor with known range [-1.0, 1.0]
    // arange needs step size: arange(start, end, step) → 100 values for 10x10
    let values: Vec<f32> = (0..100)
        .map(|i| -1.0 + (i as f32 * 0.02)) // Maps 0-99 to [-1.0, 0.98]
        .collect();
    let input = Tensor::from_vec(values, &[10, 10], &device).unwrap();

    println!("Input shape: {:?}", input.dims());
    println!("Input dtype: {:?}", input.dtype());

    // Create fake quantization layer with known scale and zero_point
    let config = QATConfig {
        quant_type: QuantizationType::Int8,
        symmetric: true,
        per_channel: false,
        calibration_batches: 10,
        fake_quant_enabled: true,
        observer_update_frequency: 10,
        ema_decay: 0.99,
    };

    // For symmetric quantization: scale = max(abs(min), abs(max)) / 127
    // With range [-1.0, 1.0], scale = 1.0 / 127 ≈ 0.00787
    let scale = 1.0 / 127.0;
    let zero_point = 127i8; // Symmetric → zero_point = 127

    let fake_quant = FakeQuantize::new(config, device.clone(), scale, zero_point).unwrap();

    println!("Scale: {}", fake_quant.scale());
    println!("Zero point: {}", fake_quant.zero_point());

    // Forward pass (quantize→dequantize)
    let output = fake_quant.forward(&input).unwrap();

    println!("Output shape: {:?}", output.dims());
    println!("Output dtype: {:?}", output.dtype());

    // Verify output shape matches input
    assert_eq!(output.dims(), input.dims(), "Output shape mismatch");
    assert_eq!(output.dtype(), input.dtype(), "Output dtype mismatch");

    // Verify quantization error is small (< 1% for this range)
    let error = output
        .sub(&input)
        .unwrap()
        .abs()
        .unwrap()
        .mean_all()
        .unwrap()
        .to_vec0::<f32>()
        .unwrap();

    println!("Quantization error (MAE): {:.6}", error);
    assert!(
        error < 0.01,
        "Quantization error too large: {} (expected < 0.01)",
        error
    );

    // Verify values are within expected range
    let output_vec = output.flatten_all().unwrap().to_vec1::<f32>().unwrap();
    for (i, &val) in output_vec.iter().enumerate() {
        assert!(
            val.is_finite(),
            "Output value at index {} is not finite: {}",
            i,
            val
        );
    }

    println!("✓ Fake quantize forward pass test PASSED");
}

// ============================================================================
// Test 2: Gradient Flow Through Fake Quantization (STE)
// ============================================================================

#[test]
fn test_fake_quantize_gradients() {
    println!("\n=== Test 2: Fake Quantize Gradient Flow ===");
    let device = test_device();
    println!("Device: {:?}", device);

    // Create test input
    let input = create_test_tensor(&device, &[8, 16]);
    println!("Input shape: {:?}", input.dims());

    // Create fake quantization layer
    let config = QATConfig::default();
    let scale = 0.1;
    let zero_point = 127i8;

    let fake_quant = FakeQuantize::new(config, device.clone(), scale, zero_point).unwrap();

    // Forward pass
    let output = fake_quant.forward(&input).unwrap();

    // Verify gradients can flow (test by checking output is differentiable)
    // In a real training loop, gradients would flow through via autograd
    // For this test, we verify the output is valid for gradient computation

    // Check output is finite (required for gradient computation)
    let output_vec = output.flatten_all().unwrap().to_vec1::<f32>().unwrap();
    for (i, &val) in output_vec.iter().enumerate() {
        assert!(
            val.is_finite(),
            "Output value at index {} is not finite: {}",
            i,
            val
        );
    }

    // Verify Straight-Through Estimator property:
    // For small perturbations, output should be close to input
    // (gradient approximation: d_output/d_input ≈ 1 for small changes)
    let perturbation = Tensor::new(&[[0.001f32]], &device).unwrap();
    let perturbed_input = input.broadcast_add(&perturbation).unwrap();
    let perturbed_output = fake_quant.forward(&perturbed_input).unwrap();

    let gradient_approx = perturbed_output
        .sub(&output)
        .unwrap()
        .mean_all()
        .unwrap()
        .to_vec0::<f32>()
        .unwrap();

    println!("Gradient approximation: {:.6}", gradient_approx);
    println!("Expected: ~0.001 (STE property)");

    // Gradient should be close to perturbation (STE: gradient flows through)
    assert!(
        (gradient_approx - 0.001).abs() < 0.01,
        "Gradient approximation incorrect: {} (expected ~0.001)",
        gradient_approx
    );

    println!("✓ Gradient flow test PASSED");
}

// ============================================================================
// Test 3: Observer Statistics Tracking (Min/Max with EMA)
// ============================================================================

#[test]
fn test_observer_statistics() {
    println!("\n=== Test 3: Observer Statistics Tracking ===");
    let device = test_device();
    println!("Device: {:?}", device);

    // Create observer with EMA decay
    let config = QATConfig {
        quant_type: QuantizationType::Int8,
        symmetric: true,
        per_channel: false,
        calibration_batches: 5,
        fake_quant_enabled: true,
        observer_update_frequency: 10,
        ema_decay: 0.9, // Fast decay for testing
    };

    let mut observer = QuantizationObserver::new(config.clone(), device.clone());

    // Initially not calibrated
    assert!(
        !observer.is_calibrated(),
        "Observer should start uncalibrated"
    );
    assert_eq!(observer.num_observations(), 0, "Should have 0 observations");

    // Feed 5 batches with known ranges
    let batch1 = Tensor::new(&[[0.0f32, 1.0f32]], &device).unwrap(); // range: [0, 1]
    let batch2 = Tensor::new(&[[-1.0f32, 2.0f32]], &device).unwrap(); // range: [-1, 2]
    let batch3 = Tensor::new(&[[-2.0f32, 1.5f32]], &device).unwrap(); // range: [-2, 1.5]
    let batch4 = Tensor::new(&[[-1.5f32, 3.0f32]], &device).unwrap(); // range: [-1.5, 3]
    let batch5 = Tensor::new(&[[-0.5f32, 2.5f32]], &device).unwrap(); // range: [-0.5, 2.5]

    println!("\nObserving 5 batches with EMA decay {}", config.ema_decay);

    observer.observe(&batch1).unwrap();
    println!("After batch 1: {:?}", observer.get_min_max());
    assert_eq!(observer.num_observations(), 1);

    observer.observe(&batch2).unwrap();
    println!("After batch 2: {:?}", observer.get_min_max());
    assert_eq!(observer.num_observations(), 2);

    observer.observe(&batch3).unwrap();
    println!("After batch 3: {:?}", observer.get_min_max());
    assert_eq!(observer.num_observations(), 3);

    observer.observe(&batch4).unwrap();
    println!("After batch 4: {:?}", observer.get_min_max());
    assert_eq!(observer.num_observations(), 4);

    observer.observe(&batch5).unwrap();
    println!("After batch 5: {:?}", observer.get_min_max());
    assert_eq!(observer.num_observations(), 5);

    // After 5 batches, should be calibrated
    assert!(
        observer.is_calibrated(),
        "Observer should be calibrated after 5 batches"
    );

    // Verify min/max are within expected range (EMA smooths extremes)
    let (min_val, max_val) = observer.get_min_max().unwrap();
    println!("\nFinal statistics:");
    println!("Min: {:.4}", min_val);
    println!("Max: {:.4}", max_val);

    // Min should be between -2.0 (extreme) and 0.0 (first batch)
    assert!(
        min_val >= -2.0 && min_val <= 0.0,
        "Min value out of expected range: {}",
        min_val
    );

    // Max should be between 1.0 (first batch) and 3.0 (extreme)
    assert!(
        max_val >= 1.0 && max_val <= 3.0,
        "Max value out of expected range: {}",
        max_val
    );

    // Test reset functionality
    observer.reset();
    assert!(
        !observer.is_calibrated(),
        "Observer should be uncalibrated after reset"
    );
    assert_eq!(
        observer.num_observations(),
        0,
        "Observations should be 0 after reset"
    );
    assert_eq!(
        observer.get_min_max(),
        None,
        "Min/max should be None after reset"
    );

    println!("✓ Observer statistics test PASSED");
}

// ============================================================================
// Test 4: QAT Calibration Phase Workflow
// ============================================================================

#[test]
fn test_qat_calibration_phase() {
    println!("\n=== Test 4: QAT Calibration Phase Workflow ===");
    let device = test_device();
    println!("Device: {:?}", device);

    // Step 1: Create calibration data (10 batches)
    let calibration_data = create_calibration_data(&device, 10, &[4, 8]);
    println!("Created {} calibration batches", calibration_data.len());

    // Step 2: Create observer
    let config = QATConfig {
        quant_type: QuantizationType::Int8,
        symmetric: true,
        per_channel: false,
        calibration_batches: 10,
        fake_quant_enabled: true,
        observer_update_frequency: 10,
        ema_decay: 0.99,
    };

    let mut observer = QuantizationObserver::new(config.clone(), device.clone());
    println!(
        "Created observer with calibration_batches: {}",
        config.calibration_batches
    );

    // Step 3: Calibration loop
    println!("\nRunning calibration loop...");
    for (i, batch) in calibration_data.iter().enumerate() {
        observer.observe(batch).unwrap();
        println!(
            "Batch {}: observations={}, calibrated={}",
            i + 1,
            observer.num_observations(),
            observer.is_calibrated()
        );
    }

    // Verify calibration complete
    assert!(
        observer.is_calibrated(),
        "Observer should be calibrated after 10 batches"
    );
    assert_eq!(
        observer.num_observations(),
        10,
        "Should have 10 observations"
    );

    // Step 4: Create FakeQuantize from observer
    let fake_quant = FakeQuantize::from_observer(&observer).unwrap();
    println!("\nCreated FakeQuantize from observer:");
    println!("Scale: {}", fake_quant.scale());
    println!("Zero point: {}", fake_quant.zero_point());

    // Verify scale and zero_point are reasonable
    assert!(fake_quant.scale() > 0.0, "Scale should be positive");
    assert!(
        fake_quant.zero_point() >= -128 && fake_quant.zero_point() <= 127,
        "Zero point should be in [-128, 127] for i8"
    );

    // Step 5: Test forward pass with calibrated fake quantization
    let test_input = create_test_tensor(&device, &[2, 8]);
    let output = fake_quant.forward(&test_input).unwrap();

    println!("\nForward pass with calibrated FakeQuantize:");
    println!("Input shape: {:?}", test_input.dims());
    println!("Output shape: {:?}", output.dims());

    assert_eq!(
        output.dims(),
        test_input.dims(),
        "Output shape should match input"
    );

    println!("✓ QAT calibration phase test PASSED");
}

// ============================================================================
// Test 5: QAT→INT8 Conversion for Deployment
// ============================================================================

#[test]
fn test_qat_to_quantized_conversion() {
    println!("\n=== Test 5: QAT→INT8 Conversion ===");
    let device = test_device();
    println!("Device: {:?}", device);

    // Create and calibrate observer
    let config = QATConfig::default();
    let mut observer = QuantizationObserver::new(config.clone(), device.clone());

    let calibration_data = create_calibration_data(&device, 100, &[16, 16]);
    for batch in &calibration_data {
        observer.observe(batch).unwrap();
    }

    assert!(observer.is_calibrated());
    println!(
        "Observer calibrated with {} batches",
        observer.num_observations()
    );

    // Create FakeQuantize from observer
    let fake_quant = FakeQuantize::from_observer(&observer).unwrap();
    println!(
        "FakeQuantize created: scale={}, zero_point={}",
        fake_quant.scale(),
        fake_quant.zero_point()
    );

    // Simulate trained weights (after QAT training)
    let trained_weights = create_test_tensor(&device, &[32, 16]);
    println!("Trained weights shape: {:?}", trained_weights.dims());

    // Convert to INT8 for deployment
    let quantized_weights = fake_quant.to_quantized(&trained_weights).unwrap();

    println!("\nQuantized weights:");
    println!("Data dtype: {:?}", quantized_weights.data.dtype());
    println!("Quantization type: {:?}", quantized_weights.quant_type);
    println!("Scale: {}", quantized_weights.scale);
    println!("Zero point: {}", quantized_weights.zero_point);

    // Verify quantized weights
    assert_eq!(
        quantized_weights.data.dtype(),
        DType::U8,
        "Quantized data should be U8"
    );
    assert_eq!(
        quantized_weights.quant_type,
        QuantizationType::Int8,
        "Should be INT8 quantization"
    );
    assert_eq!(
        quantized_weights.data.dims(),
        trained_weights.dims(),
        "Shape should be preserved"
    );

    // Verify values are valid u8 (quantized weights are stored as u8)
    let quantized_vec = quantized_weights
        .data
        .flatten_all()
        .unwrap()
        .to_vec1::<u8>()
        .unwrap();
    // All u8 values are valid by definition (0-255 range), just verify we can read them
    assert!(
        quantized_vec.len() > 0,
        "Quantized weights should have data"
    );

    // Calculate memory savings
    let original_bytes = trained_weights.dims().iter().product::<usize>() * 4; // F32 = 4 bytes
    let quantized_bytes = quantized_weights.memory_bytes();
    let savings_percent = (1.0 - (quantized_bytes as f64 / original_bytes as f64)) * 100.0;

    println!("\nMemory savings:");
    println!(
        "Original: {} bytes ({} KB)",
        original_bytes,
        original_bytes / 1024
    );
    println!(
        "Quantized: {} bytes ({} KB)",
        quantized_bytes,
        quantized_bytes / 1024
    );
    println!("Savings: {:.1}%", savings_percent);

    assert!(
        savings_percent >= 70.0,
        "Expected at least 70% memory savings, got {:.1}%",
        savings_percent
    );

    println!("✓ QAT→INT8 conversion test PASSED");
}

// ============================================================================
// Test 6: QAT vs PTQ Accuracy Comparison
// ============================================================================

#[test]
fn test_qat_accuracy_vs_ptq() {
    println!("\n=== Test 6: QAT vs PTQ Accuracy Comparison ===");
    let device = test_device();
    println!("Device: {:?}", device);

    // Create synthetic ground truth
    let ground_truth = Tensor::randn(0.0f32, 1.0f32, &[32, 10], &device).unwrap();
    println!("Ground truth shape: {:?}", ground_truth.dims());

    // Simulate FP32 model predictions (perfect accuracy for this test)
    let fp32_predictions = ground_truth.clone();

    // =======================================================================
    // PTQ Path: Direct quantization without training
    // =======================================================================
    println!("\n--- PTQ (Post-Training Quantization) ---");

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

    let mut ptq_quantizer = Quantizer::new(ptq_config, device.clone());

    // Quantize FP32 predictions directly (no training)
    let quantized_ptq = ptq_quantizer
        .quantize_tensor(&fp32_predictions, "ptq_weights")
        .unwrap();
    println!(
        "PTQ quantization: scale={}, zero_point={}",
        quantized_ptq.scale, quantized_ptq.zero_point
    );

    // Dequantize for inference
    let ptq_predictions = ptq_quantizer.dequantize_tensor(&quantized_ptq).unwrap();

    // Calculate PTQ error
    let ptq_error = ptq_predictions
        .sub(&ground_truth)
        .unwrap()
        .abs()
        .unwrap()
        .mean_all()
        .unwrap()
        .to_vec0::<f32>()
        .unwrap();
    println!("PTQ error (MAE): {:.6}", ptq_error);

    // =======================================================================
    // QAT Path: Calibration + training-aware quantization
    // =======================================================================
    println!("\n--- QAT (Quantization-Aware Training) ---");

    let qat_config = QATConfig {
        quant_type: QuantizationType::Int8,
        symmetric: true,
        per_channel: false,
        calibration_batches: 10,
        fake_quant_enabled: true,
        observer_update_frequency: 10,
        ema_decay: 0.99,
    };

    // Step 1: Calibration phase
    let mut observer = QuantizationObserver::new(qat_config.clone(), device.clone());
    let calibration_data = create_calibration_data(&device, 10, &[32, 10]);

    for batch in &calibration_data {
        observer.observe(batch).unwrap();
    }
    println!(
        "Calibration complete: {} batches observed",
        observer.num_observations()
    );

    // Step 2: Create FakeQuantize
    let fake_quant = FakeQuantize::from_observer(&observer).unwrap();
    println!(
        "FakeQuantize created: scale={}, zero_point={}",
        fake_quant.scale(),
        fake_quant.zero_point()
    );

    // Step 3: Simulate QAT training (forward pass with fake quantization)
    // In real training, this would include backprop and weight updates
    let qat_predictions = fake_quant.forward(&fp32_predictions).unwrap();

    // Calculate QAT error
    let qat_error = qat_predictions
        .sub(&ground_truth)
        .unwrap()
        .abs()
        .unwrap()
        .mean_all()
        .unwrap()
        .to_vec0::<f32>()
        .unwrap();
    println!("QAT error (MAE): {:.6}", qat_error);

    // =======================================================================
    // Comparison: QAT should be 1-2% better than PTQ
    // =======================================================================
    println!("\n--- Accuracy Comparison ---");

    let (qat_accuracy, ptq_accuracy, improvement_pct) =
        compare_qat_vs_ptq_accuracy(&qat_predictions, &ptq_predictions, &ground_truth).unwrap();

    println!("QAT accuracy: {:.4}", qat_accuracy);
    println!("PTQ accuracy: {:.4}", ptq_accuracy);
    println!("Improvement: {:.2}%", improvement_pct);

    // QAT should be at least as good as PTQ (ideally 1-2% better)
    assert!(
        qat_accuracy >= ptq_accuracy,
        "QAT accuracy ({:.4}) should be >= PTQ accuracy ({:.4})",
        qat_accuracy,
        ptq_accuracy
    );

    // For this synthetic test, we expect QAT to be slightly better
    // (in practice, QAT is 1-2% better on real datasets after full training)
    println!(
        "\nNote: QAT is {:.2}% {} than PTQ",
        improvement_pct.abs(),
        if improvement_pct >= 0.0 {
            "better"
        } else {
            "worse"
        }
    );

    println!("✓ QAT vs PTQ accuracy comparison test PASSED");
}

// ============================================================================
// Additional Edge Case Tests
// ============================================================================

#[test]
fn test_observer_error_before_calibration() {
    println!("\n=== Edge Case: Create FakeQuantize Before Calibration ===");
    let device = test_device();

    let config = QATConfig::default();
    let observer = QuantizationObserver::new(config, device);

    // Try to create FakeQuantize before calibration
    let result = FakeQuantize::from_observer(&observer);

    assert!(
        result.is_err(),
        "Should fail to create FakeQuantize before calibration"
    );
    println!("✓ Correctly rejects uncalibrated observer");
}

#[test]
fn test_fake_quantize_eval_mode() {
    println!("\n=== Edge Case: Fake Quantize Eval Mode ===");
    let device = test_device();

    let config = QATConfig::default();
    let scale = 0.1;
    let zero_point = 127i8;

    let mut fake_quant = FakeQuantize::new(config, device.clone(), scale, zero_point).unwrap();

    let input = create_test_tensor(&device, &[4, 8]);

    // Training mode: quantization applied
    fake_quant.train();
    let train_output = fake_quant.forward(&input).unwrap();

    // Eval mode: no quantization (bypass)
    fake_quant.eval();
    let eval_output = fake_quant.forward(&input).unwrap();

    // In eval mode, output should equal input (no quantization)
    let error = eval_output
        .sub(&input)
        .unwrap()
        .abs()
        .unwrap()
        .mean_all()
        .unwrap()
        .to_vec0::<f32>()
        .unwrap();

    println!("Eval mode error: {:.6} (should be ~0)", error);
    assert!(
        error < 1e-6,
        "Eval mode should bypass quantization, error: {}",
        error
    );

    // Training mode output should differ from input (quantization applied)
    let train_error = train_output
        .sub(&input)
        .unwrap()
        .abs()
        .unwrap()
        .mean_all()
        .unwrap()
        .to_vec0::<f32>()
        .unwrap();

    println!("Train mode error: {:.6} (should be > 0)", train_error);
    assert!(
        train_error > 1e-6,
        "Train mode should apply quantization, error: {}",
        train_error
    );

    println!("✓ Eval mode bypass test PASSED");
}