foxhunt/ml/tests/qat_tft_integration_test.rs

//! Integration tests for QAT-enabled TFT
//!
//! Tests the full quantization-aware training workflow:
//! 1. Create QAT wrapper from FP32 model
//! 2. Calibrate on representative data
//! 3. Run forward passes with fake quantization
//! 4. Convert to fully quantized INT8 model
//! 5. Validate accuracy preservation

use candle_core::{DType, Device, Tensor};
use ml::tft::{QATTemporalFusionTransformer, TFTConfig, TemporalFusionTransformer};
use ml::MLError;

#[test]
fn test_qat_wrapper_creation() -> Result<(), MLError> {
    // Create FP32 TFT model
    let config = TFTConfig {
        input_dim: 30,
        num_static_features: 5,
        num_known_features: 10,
        num_unknown_features: 15,
        hidden_dim: 64,
        sequence_length: 20,
        prediction_horizon: 5,
        num_quantiles: 3,
        ..Default::default()
    };

    let device = Device::Cpu;
    let fp32_model = TemporalFusionTransformer::new_with_device(config.clone(), device.clone())?;

    // Wrap with QAT
    let qat_model = QATTemporalFusionTransformer::new_from_fp32(fp32_model)?;

    // Verify QAT wrapper initialized correctly
    assert!(
        qat_model.is_calibration_mode(),
        "Should start in calibration mode"
    );
    assert_eq!(
        qat_model.num_observers(),
        1,
        "Should have 1 observer (output layer)"
    );

    Ok(())
}

#[test]
fn test_qat_forward_pass() -> Result<(), MLError> {
    // Create FP32 TFT model
    let config = TFTConfig {
        input_dim: 30,
        num_static_features: 5,
        num_known_features: 10,
        num_unknown_features: 15,
        hidden_dim: 64,
        sequence_length: 20,
        prediction_horizon: 5,
        num_quantiles: 3,
        ..Default::default()
    };

    let device = Device::Cpu;
    let fp32_model = TemporalFusionTransformer::new_with_device(config.clone(), device.clone())?;
    let mut qat_model = QATTemporalFusionTransformer::new_from_fp32(fp32_model)?;

    // Create test inputs
    let batch_size = 2;
    let static_features = Tensor::zeros(
        (batch_size, config.num_static_features),
        DType::F32,
        &device,
    )?;
    let historical_features = Tensor::zeros(
        (
            batch_size,
            config.sequence_length,
            config.num_unknown_features,
        ),
        DType::F32,
        &device,
    )?;
    let future_features = Tensor::zeros(
        (
            batch_size,
            config.prediction_horizon,
            config.num_known_features,
        ),
        DType::F32,
        &device,
    )?;

    // Forward pass should work
    let output = qat_model.forward(&static_features, &historical_features, &future_features)?;

    // Validate output shape
    let output_dims = output.dims();
    assert_eq!(output_dims.len(), 3, "Output should be 3D");
    assert_eq!(output_dims[0], batch_size, "Batch size should match");
    assert_eq!(
        output_dims[1], config.prediction_horizon,
        "Horizon should match"
    );
    assert_eq!(
        output_dims[2], config.num_quantiles,
        "Quantiles should match"
    );

    Ok(())
}

#[test]
fn test_qat_calibration() -> Result<(), MLError> {
    // Create FP32 TFT model
    let config = TFTConfig {
        input_dim: 30,
        num_static_features: 5,
        num_known_features: 10,
        num_unknown_features: 15,
        hidden_dim: 64,
        sequence_length: 20,
        prediction_horizon: 5,
        num_quantiles: 3,
        ..Default::default()
    };

    let device = Device::Cpu;
    let fp32_model = TemporalFusionTransformer::new_with_device(config.clone(), device.clone())?;
    let mut qat_model = QATTemporalFusionTransformer::new_from_fp32(fp32_model)?;

    // Generate calibration data (10 samples)
    let mut calibration_data = Vec::new();
    for _ in 0..10 {
        let static_feat = Tensor::randn(0.0f32, 1.0, (1, config.num_static_features), &device)?;
        let hist_feat = Tensor::randn(
            0.0f32,
            1.0,
            (1, config.sequence_length, config.num_unknown_features),
            &device,
        )?;
        let fut_feat = Tensor::randn(
            0.0f32,
            1.0,
            (1, config.prediction_horizon, config.num_known_features),
            &device,
        )?;
        calibration_data.push((static_feat, hist_feat, fut_feat));
    }

    // Calibrate
    qat_model.calibrate(&calibration_data)?;

    // Verify calibration disabled
    assert!(
        !qat_model.is_calibration_mode(),
        "Calibration mode should be disabled after calibration"
    );

    // Get calibration stats
    let stats = qat_model.get_calibration_stats();
    assert!(!stats.is_empty(), "Should have calibration statistics");

    // Verify all observers have positive scales and samples
    for (name, (scale, _zero_point, num_samples)) in stats {
        assert!(
            scale > 0.0,
            "Layer {} should have positive scale, got {}",
            name,
            scale
        );
        assert!(
            num_samples > 0,
            "Layer {} should have samples, got {}",
            name,
            num_samples
        );
    }

    Ok(())
}

#[test]
fn test_qat_to_quantized_conversion() -> Result<(), MLError> {
    // Create FP32 TFT model
    let config = TFTConfig {
        input_dim: 30,
        num_static_features: 5,
        num_known_features: 10,
        num_unknown_features: 15,
        hidden_dim: 64,
        sequence_length: 20,
        prediction_horizon: 5,
        num_quantiles: 3,
        ..Default::default()
    };

    let device = Device::Cpu;
    let fp32_model = TemporalFusionTransformer::new_with_device(config.clone(), device.clone())?;
    let mut qat_model = QATTemporalFusionTransformer::new_from_fp32(fp32_model)?;

    // Generate calibration data
    let mut calibration_data = Vec::new();
    for _ in 0..10 {
        let static_feat = Tensor::randn(0.0f32, 1.0, (1, config.num_static_features), &device)?;
        let hist_feat = Tensor::randn(
            0.0f32,
            1.0,
            (1, config.sequence_length, config.num_unknown_features),
            &device,
        )?;
        let fut_feat = Tensor::randn(
            0.0f32,
            1.0,
            (1, config.prediction_horizon, config.num_known_features),
            &device,
        )?;
        calibration_data.push((static_feat, hist_feat, fut_feat));
    }

    // Calibrate
    qat_model.calibrate(&calibration_data)?;

    // Convert to fully quantized INT8 model
    let int8_model = qat_model.to_quantized()?;

    // Verify INT8 model created
    assert_eq!(
        int8_model.config.input_dim, config.input_dim,
        "INT8 model should have same config"
    );

    // Verify memory reduction
    let int8_memory = int8_model.memory_usage_bytes();
    let fp32_memory = 125 * 1024 * 1024; // 125MB base
    let reduction_ratio = (int8_memory as f64) / (fp32_memory as f64);

    assert!(
        reduction_ratio < 0.5,
        "INT8 model should have <50% memory vs FP32, got {:.2}%",
        reduction_ratio * 100.0
    );

    Ok(())
}

#[test]
fn test_qat_memory_usage() -> Result<(), MLError> {
    // Create FP32 TFT model
    let config = TFTConfig {
        input_dim: 30,
        num_static_features: 5,
        num_known_features: 10,
        num_unknown_features: 15,
        hidden_dim: 64,
        ..Default::default()
    };

    let device = Device::Cpu;
    let fp32_model = TemporalFusionTransformer::new_with_device(config.clone(), device.clone())?;
    let qat_model = QATTemporalFusionTransformer::new_from_fp32(fp32_model)?;

    let memory = qat_model.memory_usage();

    // Should be approximately FP32 model size (~125MB + small observer overhead)
    assert!(
        memory > 125 * 1024 * 1024,
        "Memory should be at least FP32 baseline"
    );
    assert!(
        memory < 130 * 1024 * 1024,
        "Memory should have small observer overhead"
    );

    Ok(())
}

#[test]
fn test_qat_uncalibrated_conversion_fails() -> Result<(), MLError> {
    // Create FP32 TFT model
    let config = TFTConfig {
        input_dim: 30,
        num_static_features: 5,
        num_known_features: 10,
        num_unknown_features: 15,
        hidden_dim: 64,
        ..Default::default()
    };

    let device = Device::Cpu;
    let fp32_model = TemporalFusionTransformer::new_with_device(config.clone(), device.clone())?;
    let qat_model = QATTemporalFusionTransformer::new_from_fp32(fp32_model)?;

    // Try to convert without calibration
    let result = qat_model.to_quantized();

    // Should fail because observers not calibrated
    assert!(
        result.is_err(),
        "Conversion should fail without calibration"
    );

    if let Err(e) = result {
        let err_msg = format!("{:?}", e);
        assert!(
            err_msg.contains("not calibrated"),
            "Error should mention calibration requirement"
        );
    }

    Ok(())
}

#[test]
fn test_fake_quantize_statistics_collection() -> Result<(), MLError> {
    use ml::memory_optimization::{FakeQuantize, QATConfig, QuantizationObserver};

    let device = Device::Cpu;
    let config = QATConfig {
        calibration_batches: 2,
        ..Default::default()
    };
    let mut observer = QuantizationObserver::new(config.clone(), device.clone());

    // Create test tensor with known range [-2.0, 3.0]
    let x = Tensor::new(&[[-2.0f32, -1.0, 0.0, 1.0, 2.0, 3.0]], &device)?;

    // Run calibration
    observer.observe(&x)?;

    // Should have running statistics after first observation
    let stats = observer.get_min_max();
    assert!(stats.is_some(), "Should have statistics after observation");

    // Run another sample
    let x2 = Tensor::new(&[[-1.5f32, -0.5, 0.5, 1.5, 2.5]], &device)?;
    observer.observe(&x2)?;

    // Check calibration complete
    assert!(
        observer.is_calibrated(),
        "Should be calibrated after 2 batches"
    );

    // Create FakeQuantize from calibrated observer
    let fake_quant = FakeQuantize::from_observer(&observer)?;

    // Should have frozen scale/zero_point
    let (scale, zero_point) = fake_quant.scale_zero_point();

    // Verify symmetric quantization parameters
    assert!(scale > 0.0, "Scale should be positive");
    assert_eq!(
        zero_point, 127,
        "Symmetric quantization uses zero_point=127"
    );

    // Verify scale is reasonable for the input range
    // abs_max = max(2.0, 3.0) = 3.0
    // scale = 3.0 / 127 ≈ 0.024
    let expected_scale = 3.0 / 127.0;
    let scale_tolerance = 0.01; // Allow some EMA variance
    assert!(
        (scale - expected_scale).abs() < scale_tolerance,
        "Scale should be approximately {:.6}, got {:.6}",
        expected_scale,
        scale
    );

    Ok(())
}

#[test]
fn test_fake_quantize_noise_simulation() -> Result<(), MLError> {
    use ml::memory_optimization::{FakeQuantize, QATConfig, QuantizationObserver};

    let device = Device::Cpu;
    let config = QATConfig {
        calibration_batches: 1,
        ..Default::default()
    };
    let mut observer = QuantizationObserver::new(config.clone(), device.clone());

    // Create test tensor
    let x = Tensor::new(&[[1.0f32, 2.0, 3.0, 4.0]], &device)?;

    // Calibrate on this tensor
    observer.observe(&x)?;
    let fake_quant = FakeQuantize::from_observer(&observer)?;

    // Run quantization
    let quantized = fake_quant.forward(&x)?;

    // Verify output is not exactly equal to input (quantization noise)
    let x_vec = x.flatten_all()?.to_vec1::<f32>()?;
    let q_vec = quantized.flatten_all()?.to_vec1::<f32>()?;

    // Check that values are different but close
    let mut has_noise = false;
    for (orig, quant) in x_vec.iter().zip(q_vec.iter()) {
        let diff = (orig - quant).abs();
        if diff > 1e-6 {
            has_noise = true;
        }
        // Values should be within quantization error
        assert!(
            diff < 0.1,
            "Quantization error too large: original={}, quantized={}",
            orig,
            quant
        );
    }

    assert!(
        has_noise,
        "Fake quantization should introduce some quantization noise"
    );

    Ok(())
}

#[test]
fn test_qat_end_to_end_workflow() -> Result<(), MLError> {
    // 1. Create FP32 TFT model
    let config = TFTConfig {
        input_dim: 30,
        num_static_features: 5,
        num_known_features: 10,
        num_unknown_features: 15,
        hidden_dim: 64,
        sequence_length: 20,
        prediction_horizon: 5,
        num_quantiles: 3,
        ..Default::default()
    };

    let device = Device::Cpu;
    let fp32_model = TemporalFusionTransformer::new_with_device(config.clone(), device.clone())?;

    // 2. Wrap with QAT
    let mut qat_model = QATTemporalFusionTransformer::new_from_fp32(fp32_model)?;

    // 3. Generate calibration data
    let mut calibration_data = Vec::new();
    for _ in 0..20 {
        let static_feat = Tensor::randn(0.0f32, 1.0, (1, config.num_static_features), &device)?;
        let hist_feat = Tensor::randn(
            0.0f32,
            1.0,
            (1, config.sequence_length, config.num_unknown_features),
            &device,
        )?;
        let fut_feat = Tensor::randn(
            0.0f32,
            1.0,
            (1, config.prediction_horizon, config.num_known_features),
            &device,
        )?;
        calibration_data.push((static_feat, hist_feat, fut_feat));
    }

    // 4. Calibrate
    qat_model.calibrate(&calibration_data)?;

    // 5. Get calibration stats
    let stats = qat_model.get_calibration_stats();
    println!("Calibration stats: {} observers calibrated", stats.len());
    for (name, (scale, zero_point, num_samples)) in &stats {
        println!(
            "  {}: scale={:.6}, zero_point={}, samples={}",
            name, scale, zero_point, num_samples
        );
    }

    // 6. Run inference with fake quantization
    let static_feat = Tensor::randn(0.0f32, 1.0, (2, config.num_static_features), &device)?;
    let hist_feat = Tensor::randn(
        0.0f32,
        1.0,
        (2, config.sequence_length, config.num_unknown_features),
        &device,
    )?;
    let fut_feat = Tensor::randn(
        0.0f32,
        1.0,
        (2, config.prediction_horizon, config.num_known_features),
        &device,
    )?;

    let qat_output = qat_model.forward(&static_feat, &hist_feat, &fut_feat)?;

    // 7. Convert to fully quantized INT8
    let int8_model = qat_model.to_quantized()?;

    // 8. Verify memory reduction
    let int8_memory = int8_model.memory_usage_bytes();
    let fp32_memory = 125 * 1024 * 1024;
    let reduction = (1.0 - (int8_memory as f64 / fp32_memory as f64)) * 100.0;
    println!(
        "Memory reduction: {:.1}% (FP32: {}MB, INT8: {}MB)",
        reduction,
        fp32_memory / (1024 * 1024),
        int8_memory / (1024 * 1024)
    );

    // 9. Verify output shapes match
    assert_eq!(
        qat_output.dims(),
        &[2, config.prediction_horizon, config.num_quantiles]
    );

    Ok(())
}