foxhunt/ml/tests/tft_complete_int8_integration_test.rs

//! Complete INT8 TFT Integration Test
//!
//! Validates end-to-end quantization of TFT model:
//! - Load F32 TFT model
//! - Convert to INT8 (VSN, LSTM, Attention, GRN)
//! - Verify forward pass integrity
//! - Validate accuracy loss <5%
//! - Verify memory reduction 70-80%
//! - Validate checkpoint save/load
//!
//! Target: 2,952MB → 738MB (75% reduction)

use anyhow::Result;
use candle_core::{DType, Device, Tensor};
use candle_nn::{VarBuilder, VarMap};
use std::sync::Arc;

use ml::memory_optimization::quantization::{
    QuantizationConfig, QuantizationType,
};
use ml::tft::{TemporalFusionTransformer, TFTConfig};
use ml::MLError;

// Import quantized TFT (to be implemented)
use ml::tft::quantized_tft::QuantizedTFT;

/// Helper: Create small TFT model for testing
fn create_test_tft() -> Result<TemporalFusionTransformer> {
    let config = TFTConfig {
        input_dim: 32,
        hidden_dim: 64,
        num_heads: 4,
        num_layers: 2,
        prediction_horizon: 5,
        sequence_length: 10,
        num_quantiles: 5,
        num_static_features: 4,
        num_known_features: 8,
        num_unknown_features: 16,
        learning_rate: 1e-3,
        batch_size: 32,
        dropout_rate: 0.1,
        l2_regularization: 1e-4,
        use_flash_attention: false,
        mixed_precision: false,
        memory_efficient: true,
        max_inference_latency_us: 50,
        target_throughput_pps: 100_000,
    };

    TemporalFusionTransformer::new(config)
        .map_err(|e| anyhow::anyhow!("Failed to create TFT: {:?}", e))
}

/// Helper: Generate random test inputs
fn generate_test_inputs(
    config: &TFTConfig,
    batch_size: usize,
    device: &Device,
) -> Result<(Tensor, Tensor, Tensor)> {
    let static_features = Tensor::randn(
        0.0f32,
        1.0f32,
        (batch_size, config.num_static_features),
        device,
    )?;

    let historical_features = Tensor::randn(
        0.0f32,
        1.0f32,
        (batch_size, config.sequence_length, config.num_unknown_features),
        device,
    )?;

    let future_features = Tensor::randn(
        0.0f32,
        1.0f32,
        (batch_size, config.prediction_horizon, config.num_known_features),
        device,
    )?;

    Ok((static_features, historical_features, future_features))
}

/// Helper: Calculate relative error between F32 and INT8 predictions
fn calculate_relative_error(f32_pred: &Tensor, int8_pred: &Tensor) -> Result<f64> {
    let diff = (f32_pred - int8_pred)?.abs()?;
    let abs_f32 = f32_pred.abs()?;
    let relative_error = (&diff / &abs_f32)?;
    let mean_error = relative_error.mean_all()?.to_vec0::<f32>()?;
    Ok(mean_error as f64)
}

/// Helper: Estimate model memory size (rough approximation)
fn estimate_model_memory_mb(varmap: &VarMap) -> Result<f64> {
    let var_data = varmap.data().lock().unwrap();
    let mut total_bytes = 0usize;

    for (_name, tensor) in var_data.iter() {
        let elem_count = tensor.elem_count();
        let dtype = tensor.dtype();
        let bytes_per_elem = match dtype {
            DType::F32 => 4,
            DType::F64 => 8,
            DType::U8 => 1,
            DType::I64 => 8,
            _ => 4, // default assumption
        };
        total_bytes += elem_count * bytes_per_elem;
    }

    Ok(total_bytes as f64 / (1024.0 * 1024.0))
}

// ============================================================================
// Test 1: Load F32 TFT and convert to INT8
// ============================================================================

#[test]
fn test_f32_to_int8_conversion() -> Result<()> {
    println!("\n=== Test 1: F32 → INT8 Conversion ===");

    // 1. Create F32 TFT model
    let f32_tft = create_test_tft()?;
    println!("✓ Created F32 TFT model");

    // 2. Create quantization config
    let quant_config = QuantizationConfig {
        quant_type: QuantizationType::PerChannel,
        calibration_method: ml::memory_optimization::quantization::CalibrationMethod::MinMax,
        bits: 8,
    };
    println!("✓ Created quantization config: {:?}", quant_config);

    // 3. Convert to INT8
    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);
    let int8_tft = QuantizedTFT::from_f32_model(&f32_tft, quant_config, device.clone())?;
    println!("✓ Converted to INT8 TFT");

    // 4. Verify dimensions match
    assert_eq!(int8_tft.config().input_dim, f32_tft.config.input_dim);
    assert_eq!(int8_tft.config().hidden_dim, f32_tft.config.hidden_dim);
    assert_eq!(int8_tft.config().num_heads, f32_tft.config.num_heads);
    println!("✓ Dimensions match");

    // 5. Verify quantized components exist
    assert!(int8_tft.has_quantized_vsn(), "Missing quantized VSN");
    assert!(int8_tft.has_quantized_lstm(), "Missing quantized LSTM");
    assert!(int8_tft.has_quantized_attention(), "Missing quantized Attention");
    assert!(int8_tft.has_quantized_grn(), "Missing quantized GRN");
    println!("✓ All quantized components present");

    Ok(())
}

// ============================================================================
// Test 2: Forward pass end-to-end
// ============================================================================

#[test]
fn test_quantized_forward_pass() -> Result<()> {
    println!("\n=== Test 2: Quantized Forward Pass ===");

    // 1. Create models
    let mut f32_tft = create_test_tft()?;
    let quant_config = QuantizationConfig {
        quant_type: QuantizationType::PerChannel,
        calibration_method: ml::memory_optimization::quantization::CalibrationMethod::MinMax,
        bits: 8,
    };
    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);
    let mut int8_tft = QuantizedTFT::from_f32_model(&f32_tft, quant_config, device.clone())?;
    println!("✓ Created F32 and INT8 models");

    // 2. Generate test inputs
    let batch_size = 4;
    let (static_features, historical_features, future_features) =
        generate_test_inputs(&f32_tft.config, batch_size, &device)?;
    println!("✓ Generated test inputs (batch_size={})", batch_size);

    // 3. F32 forward pass
    let f32_output = f32_tft.forward(&static_features, &historical_features, &future_features)?;
    let f32_shape = f32_output.dims();
    println!("✓ F32 forward pass: shape={:?}", f32_shape);

    // 4. INT8 forward pass
    let int8_output = int8_tft.forward(&static_features, &historical_features, &future_features)?;
    let int8_shape = int8_output.dims();
    println!("✓ INT8 forward pass: shape={:?}", int8_shape);

    // 5. Verify shapes match
    assert_eq!(
        f32_shape, int8_shape,
        "Output shapes mismatch: F32={:?} vs INT8={:?}",
        f32_shape, int8_shape
    );
    println!("✓ Output shapes match");

    // 6. Verify no NaN/Inf
    let int8_data = int8_output.flatten_all()?.to_vec1::<f32>()?;
    let has_nan = int8_data.iter().any(|x| x.is_nan());
    let has_inf = int8_data.iter().any(|x| x.is_infinite());
    assert!(!has_nan, "INT8 output contains NaN");
    assert!(!has_inf, "INT8 output contains Inf");
    println!("✓ No NaN/Inf in output");

    Ok(())
}

// ============================================================================
// Test 3: Accuracy loss <5%
// ============================================================================

#[test]
fn test_accuracy_loss_under_5_percent() -> Result<()> {
    println!("\n=== Test 3: Accuracy Loss <5% ===");

    // 1. Create models
    let mut f32_tft = create_test_tft()?;
    let quant_config = QuantizationConfig {
        quant_type: QuantizationType::PerChannel,
        calibration_method: ml::memory_optimization::quantization::CalibrationMethod::MinMax,
        bits: 8,
    };
    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);
    let mut int8_tft = QuantizedTFT::from_f32_model(&f32_tft, quant_config, device.clone())?;
    println!("✓ Created models");

    // 2. Run multiple forward passes to get average error
    let num_samples = 10;
    let mut total_error = 0.0;

    for i in 0..num_samples {
        let (static_features, historical_features, future_features) =
            generate_test_inputs(&f32_tft.config, 4, &device)?;

        let f32_output = f32_tft.forward(&static_features, &historical_features, &future_features)?;
        let int8_output = int8_tft.forward(&static_features, &historical_features, &future_features)?;

        let rel_error = calculate_relative_error(&f32_output, &int8_output)?;
        total_error += rel_error;

        println!("  Sample {}: relative error = {:.4}%", i + 1, rel_error * 100.0);
    }

    let avg_error = total_error / num_samples as f64;
    println!("\n✓ Average relative error: {:.4}%", avg_error * 100.0);

    // 3. Verify <5% accuracy loss
    assert!(
        avg_error < 0.05,
        "Accuracy loss {:.4}% exceeds 5% threshold",
        avg_error * 100.0
    );
    println!("✓ Accuracy loss within 5% threshold");

    Ok(())
}

// ============================================================================
// Test 4: Memory reduction 70-80%
// ============================================================================

#[test]
fn test_memory_reduction_70_to_80_percent() -> Result<()> {
    println!("\n=== Test 4: Memory Reduction 70-80% ===");

    // 1. Create F32 model and estimate memory
    let f32_tft = create_test_tft()?;
    let f32_memory_mb = estimate_model_memory_mb(&f32_tft.varmap)?;
    println!("✓ F32 model memory: {:.2} MB", f32_memory_mb);

    // 2. Convert to INT8
    let quant_config = QuantizationConfig {
        quant_type: QuantizationType::PerChannel,
        calibration_method: ml::memory_optimization::quantization::CalibrationMethod::MinMax,
        bits: 8,
    };
    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);
    let int8_tft = QuantizedTFT::from_f32_model(&f32_tft, quant_config, device)?;

    // 3. Estimate INT8 memory (including scale/zero_point overhead)
    let int8_memory_mb = int8_tft.estimate_memory_usage_mb()?;
    println!("✓ INT8 model memory: {:.2} MB", int8_memory_mb);

    // 4. Calculate reduction
    let reduction = (f32_memory_mb - int8_memory_mb) / f32_memory_mb;
    println!("✓ Memory reduction: {:.2}%", reduction * 100.0);

    // 5. Verify 70-80% reduction (allowing some overhead)
    assert!(
        reduction >= 0.65 && reduction <= 0.85,
        "Memory reduction {:.2}% not in 65-85% range (target: 70-80%)",
        reduction * 100.0
    );
    println!("✓ Memory reduction within expected range");

    Ok(())
}

// ============================================================================
// Test 5: Checkpoint save/load
// ============================================================================

#[test]
fn test_checkpoint_save_load() -> Result<()> {
    println!("\n=== Test 5: Checkpoint Save/Load ===");

    // 1. Create and convert model
    let f32_tft = create_test_tft()?;
    let quant_config = QuantizationConfig {
        quant_type: QuantizationType::PerChannel,
        calibration_method: ml::memory_optimization::quantization::CalibrationMethod::MinMax,
        bits: 8,
    };
    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);
    let mut int8_tft = QuantizedTFT::from_f32_model(&f32_tft, quant_config.clone(), device.clone())?;
    println!("✓ Created INT8 model");

    // 2. Run forward pass to get baseline output
    let (static_features, historical_features, future_features) =
        generate_test_inputs(&f32_tft.config, 4, &device)?;
    let output_before = int8_tft.forward(&static_features, &historical_features, &future_features)?;
    println!("✓ Generated baseline output");

    // 3. Save checkpoint
    let checkpoint_data = int8_tft.serialize_state()?;
    println!("✓ Serialized checkpoint: {} bytes", checkpoint_data.len());

    // 4. Create new model and load checkpoint
    let f32_tft_new = create_test_tft()?;
    let mut int8_tft_new = QuantizedTFT::from_f32_model(&f32_tft_new, quant_config, device.clone())?;
    int8_tft_new.deserialize_state(&checkpoint_data)?;
    println!("✓ Loaded checkpoint into new model");

    // 5. Run forward pass with loaded model
    let output_after = int8_tft_new.forward(&static_features, &historical_features, &future_features)?;
    println!("✓ Forward pass with loaded model");

    // 6. Verify outputs match
    let diff = (&output_before - &output_after)?.abs()?.sum_all()?.to_vec0::<f32>()?;
    println!("✓ Output difference: {:.6e}", diff);

    assert!(
        diff < 1e-4,
        "Checkpoint load/save outputs differ by {:.6e}",
        diff
    );
    println!("✓ Checkpoint save/load successful");

    Ok(())
}

// ============================================================================
// Test 6: Batch processing
// ============================================================================

#[test]
fn test_batch_processing() -> Result<()> {
    println!("\n=== Test 6: Batch Processing ===");

    let f32_tft = create_test_tft()?;
    let quant_config = QuantizationConfig {
        quant_type: QuantizationType::PerChannel,
        calibration_method: ml::memory_optimization::quantization::CalibrationMethod::MinMax,
        bits: 8,
    };
    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);
    let mut int8_tft = QuantizedTFT::from_f32_model(&f32_tft, quant_config, device.clone())?;
    println!("✓ Created INT8 model");

    // Test different batch sizes
    for batch_size in [1, 4, 8, 16] {
        let (static_features, historical_features, future_features) =
            generate_test_inputs(&f32_tft.config, batch_size, &device)?;

        let output = int8_tft.forward(&static_features, &historical_features, &future_features)?;
        let output_shape = output.dims();

        assert_eq!(
            output_shape[0], batch_size,
            "Batch size mismatch: expected {} got {}",
            batch_size, output_shape[0]
        );

        println!("  ✓ Batch size {}: output shape {:?}", batch_size, output_shape);
    }

    println!("✓ All batch sizes processed successfully");
    Ok(())
}

// ============================================================================
// Test 7: Component-level quantization verification
// ============================================================================

#[test]
fn test_component_quantization() -> Result<()> {
    println!("\n=== Test 7: Component-Level Quantization ===");

    let f32_tft = create_test_tft()?;
    let quant_config = QuantizationConfig {
        quant_type: QuantizationType::PerChannel,
        calibration_method: ml::memory_optimization::quantization::CalibrationMethod::MinMax,
        bits: 8,
    };
    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);
    let int8_tft = QuantizedTFT::from_f32_model(&f32_tft, quant_config, device)?;

    // 1. Verify VSN quantization
    let vsn_quantized = int8_tft.has_quantized_vsn();
    assert!(vsn_quantized, "VSN not quantized");
    println!("  ✓ VSN quantized");

    // 2. Verify LSTM quantization
    let lstm_quantized = int8_tft.has_quantized_lstm();
    assert!(lstm_quantized, "LSTM not quantized");
    println!("  ✓ LSTM quantized");

    // 3. Verify Attention quantization
    let attention_quantized = int8_tft.has_quantized_attention();
    assert!(attention_quantized, "Attention not quantized");
    println!("  ✓ Attention quantized");

    // 4. Verify GRN quantization
    let grn_quantized = int8_tft.has_quantized_grn();
    assert!(grn_quantized, "GRN not quantized");
    println!("  ✓ GRN quantized");

    println!("✓ All components quantized successfully");
    Ok(())
}

// ============================================================================
// Test 8: Quantization dtype verification
// ============================================================================

#[test]
fn test_quantized_dtypes() -> Result<()> {
    println!("\n=== Test 8: Quantized DTypes ===");

    let f32_tft = create_test_tft()?;
    let quant_config = QuantizationConfig {
        quant_type: QuantizationType::PerChannel,
        calibration_method: ml::memory_optimization::quantization::CalibrationMethod::MinMax,
        bits: 8,
    };
    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);
    let int8_tft = QuantizedTFT::from_f32_model(&f32_tft, quant_config, device)?;

    // Verify all quantized weights use U8 dtype
    let all_u8 = int8_tft.verify_all_weights_u8()?;
    assert!(all_u8, "Not all quantized weights are U8 dtype");
    println!("✓ All quantized weights use U8 dtype");

    Ok(())
}

// ============================================================================
// Integration Test: Full pipeline with realistic config
// ============================================================================

#[test]
fn test_full_pipeline_realistic_config() -> Result<()> {
    println!("\n=== Integration Test: Realistic TFT Quantization ===");

    // 1. Create realistic TFT config (similar to production)
    let config = TFTConfig {
        input_dim: 64,
        hidden_dim: 128,
        num_heads: 8,
        num_layers: 3,
        prediction_horizon: 10,
        sequence_length: 50,
        num_quantiles: 9,
        num_static_features: 5,
        num_known_features: 10,
        num_unknown_features: 20,
        learning_rate: 1e-3,
        batch_size: 64,
        dropout_rate: 0.1,
        l2_regularization: 1e-4,
        use_flash_attention: false,
        mixed_precision: false,
        memory_efficient: true,
        max_inference_latency_us: 50,
        target_throughput_pps: 100_000,
    };

    let mut f32_tft = TemporalFusionTransformer::new(config.clone())
        .map_err(|e| anyhow::anyhow!("Failed to create TFT: {:?}", e))?;
    println!("✓ Created realistic F32 TFT");

    // 2. Quantize to INT8
    let quant_config = QuantizationConfig {
        quant_type: QuantizationType::PerChannel,
        calibration_method: ml::memory_optimization::quantization::CalibrationMethod::MinMax,
        bits: 8,
    };
    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);
    let mut int8_tft = QuantizedTFT::from_f32_model(&f32_tft, quant_config, device.clone())?;
    println!("✓ Converted to INT8");

    // 3. Run inference with realistic batch
    let batch_size = 32;
    let (static_features, historical_features, future_features) =
        generate_test_inputs(&config, batch_size, &device)?;

    let f32_output = f32_tft.forward(&static_features, &historical_features, &future_features)?;
    let int8_output = int8_tft.forward(&static_features, &historical_features, &future_features)?;
    println!("✓ Forward passes completed");

    // 4. Verify accuracy
    let rel_error = calculate_relative_error(&f32_output, &int8_output)?;
    println!("✓ Relative error: {:.4}%", rel_error * 100.0);
    assert!(
        rel_error < 0.05,
        "Accuracy loss {:.4}% exceeds 5%",
        rel_error * 100.0
    );

    // 5. Report memory savings
    let f32_memory = estimate_model_memory_mb(&f32_tft.varmap)?;
    let int8_memory = int8_tft.estimate_memory_usage_mb()?;
    let reduction = (f32_memory - int8_memory) / f32_memory;
    println!("✓ Memory: F32={:.2}MB, INT8={:.2}MB, Reduction={:.2}%",
        f32_memory, int8_memory, reduction * 100.0);

    println!("\n=== Integration Test PASSED ===");
    Ok(())
}