foxhunt/ml/tests/tft_int8_calibration_dataset_test.rs

//! TFT INT8 Calibration Dataset Tests
//!
//! Test-driven development for INT8 quantization calibration using ES.FUT data.
//! Collects activation statistics for optimal per-layer quantization.

use candle_core::{DType, Device, Tensor};
use std::path::PathBuf;

use ml::tft::{TemporalFusionTransformer, TFTConfig};
use ml::data_loaders::DbnSequenceLoader;

/// Test 1: Load 1,000 bars from ES.FUT DBN file
#[tokio::test]
async fn test_load_calibration_bars_from_es_fut() -> Result<(), Box<dyn std::error::Error>> {
    // ES.FUT file path
    let dbn_file = PathBuf::from("test_data/real/databento");

    // Skip test if file doesn't exist
    if !dbn_file.exists() {
        println!("⚠️  Skipping test: DBN directory not found");
        return Ok(());
    }

    // Load 1,000 bars for calibration
    let mut loader = DbnSequenceLoader::with_limits(60, 256, Some(100), 10).await?;
    let (train_data, _val_data) = loader.load_sequences(&dbn_file, 0.9).await?;

    // Verify we got enough data
    assert!(train_data.len() >= 50, "Need at least 50 sequences for calibration (got {})", train_data.len());
    println!("✅ Loaded {} sequences for calibration", train_data.len());

    // Verify feature dimensions [batch=1, seq_len=60, d_model=256]
    let (input, _target) = &train_data[0];
    let dims = input.dims();
    assert_eq!(dims.len(), 3, "Input should be 3D tensor");
    assert_eq!(dims[0], 1, "Batch size should be 1");
    assert_eq!(dims[1], 60, "Sequence length should be 60");
    assert_eq!(dims[2], 256, "Feature dimension should be 256");
    println!("✅ Feature dimensions correct: {:?}", dims);

    Ok(())
}

/// Test 2: Extract 256-dimensional features from OHLCV bars
#[tokio::test]
async fn test_extract_256_dim_features() -> Result<(), Box<dyn std::error::Error>> {
    // Load data
    let dbn_file = PathBuf::from("test_data/real/databento");

    if !dbn_file.exists() {
        println!("⚠️  Skipping test: DBN directory not found");
        return Ok(());
    }

    let mut loader = DbnSequenceLoader::with_limits(60, 256, Some(10), 10).await?;
    let (train_data, _val_data) = loader.load_sequences(&dbn_file, 0.9).await?;

    assert!(!train_data.is_empty(), "Need at least 1 sequence");

    // Extract features from first sequence
    let (input, _target) = &train_data[0];

    // Verify feature extraction
    let feature_vec = input.flatten_all()?.to_vec1::<f64>()?;
    assert_eq!(feature_vec.len(), 60 * 256, "Feature vector size mismatch");

    // Check for valid numerical range (normalized features should be ~[-3, 3])
    let feature_stats: (f64, f64) = feature_vec.iter().fold((f64::MAX, f64::MIN), |(min, max), &x| {
        (min.min(x), max.max(x))
    });
    println!("✅ Feature range: [{:.2}, {:.2}]", feature_stats.0, feature_stats.1);

    // Check for NaN or Inf
    assert!(feature_vec.iter().all(|x| x.is_finite()), "Features contain NaN or Inf");

    Ok(())
}

/// Test 3: Collect activation statistics from forward passes
#[tokio::test]
async fn test_collect_activation_statistics() -> Result<(), Box<dyn std::error::Error>> {
    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);

    // Create minimal TFT for calibration
    let config = TFTConfig {
        input_dim: 256,
        hidden_dim: 64,
        num_heads: 4,
        num_layers: 2,
        prediction_horizon: 10,
        sequence_length: 60,
        num_quantiles: 3,
        num_static_features: 2,
        num_known_features: 3,
        num_unknown_features: 256,
        batch_size: 1,
        ..Default::default()
    };

    let mut tft = TemporalFusionTransformer::new(config)?;

    // Load calibration data
    let dbn_file = PathBuf::from("test_data/real/databento");

    if !dbn_file.exists() {
        println!("⚠️  Skipping test: DBN directory not found");
        return Ok(());
    }

    let mut loader = DbnSequenceLoader::with_limits(60, 256, Some(10), 10).await?;
    let (train_data, _val_data) = loader.load_sequences(&dbn_file, 0.9).await?;

    assert!(!train_data.is_empty(), "Need at least 1 sequence");

    // Run forward passes and collect activation statistics
    let mut activation_mins = Vec::new();
    let mut activation_maxs = Vec::new();

    for (input, _target) in train_data.iter().take(10) {
        // Split input for TFT (static, historical, future)
        // For simplicity: static=[batch, 2], historical=[batch, 60, 256], future=[batch, 10, 3]
        let batch = input.dims()[0];

        // Create dummy static features [batch, 2]
        let static_features = Tensor::zeros((batch, 2), DType::F32, &device)?;

        // Use input as historical features [batch, 60, 256]
        let historical_features = input.to_dtype(DType::F32)?;

        // Create dummy future features [batch, 10, 3]
        let future_features = Tensor::zeros((batch, 10, 3), DType::F32, &device)?;

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

        // Collect activation statistics
        let output_vec = output.flatten_all()?.to_vec1::<f32>()?;
        let min_val = output_vec.iter().cloned().fold(f32::INFINITY, f32::min);
        let max_val = output_vec.iter().cloned().fold(f32::NEG_INFINITY, f32::max);

        activation_mins.push(min_val);
        activation_maxs.push(max_val);
    }

    // Verify we collected statistics
    assert_eq!(activation_mins.len(), 10.min(train_data.len()), "Should collect stats for all samples");
    assert_eq!(activation_maxs.len(), 10.min(train_data.len()), "Should collect stats for all samples");

    // Calculate global min/max for quantization
    let global_min = activation_mins.iter().cloned().fold(f32::INFINITY, f32::min);
    let global_max = activation_maxs.iter().cloned().fold(f32::NEG_INFINITY, f32::max);

    println!("✅ Activation range: [{:.6}, {:.6}]", global_min, global_max);
    assert!(global_min.is_finite() && global_max.is_finite(), "Invalid activation statistics");

    Ok(())
}

/// Test 4: Calculate scale and zero_point per layer
#[tokio::test]
async fn test_calculate_quantization_params_per_layer() -> Result<(), Box<dyn std::error::Error>> {
    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);

    // Create TFT
    let config = TFTConfig {
        input_dim: 256,
        hidden_dim: 64,
        num_heads: 4,
        num_layers: 2,
        prediction_horizon: 10,
        sequence_length: 60,
        num_quantiles: 3,
        num_static_features: 2,
        num_known_features: 3,
        num_unknown_features: 256,
        batch_size: 1,
        ..Default::default()
    };

    let mut tft = TemporalFusionTransformer::new(config)?;

    // Load calibration data
    let dbn_file = PathBuf::from("test_data/real/databento");

    if !dbn_file.exists() {
        println!("⚠️  Skipping test: DBN directory not found");
        return Ok(());
    }

    let mut loader = DbnSequenceLoader::with_limits(60, 256, Some(10), 10).await?;
    let (train_data, _val_data) = loader.load_sequences(&dbn_file, 0.9).await?;

    assert!(!train_data.is_empty(), "Need at least 1 sequence");

    // Simulate per-layer activation collection
    // In real implementation, this would hook into each layer's output
    let mut layer_stats = std::collections::HashMap::new();

    for (idx, (input, _target)) in train_data.iter().take(10).enumerate() {
        // Run forward pass
        let batch = input.dims()[0];
        let static_features = Tensor::zeros((batch, 2), DType::F32, &device)?;
        let historical_features = input.to_dtype(DType::F32)?;
        let future_features = Tensor::zeros((batch, 10, 3), DType::F32, &device)?;

        let output = tft.forward(&static_features, &historical_features, &future_features)?;

        // Collect stats for "output_layer" (in real implementation, hook each layer)
        let output_vec = output.flatten_all()?.to_vec1::<f32>()?;
        let min_val = output_vec.iter().cloned().fold(f32::INFINITY, f32::min);
        let max_val = output_vec.iter().cloned().fold(f32::NEG_INFINITY, f32::max);

        let entry = layer_stats.entry("output_layer".to_string()).or_insert((Vec::new(), Vec::new()));
        entry.0.push(min_val);
        entry.1.push(max_val);

        if idx == 0 {
            println!("✅ Sample {}: activation range [{:.6}, {:.6}]", idx, min_val, max_val);
        }
    }

    // Calculate quantization parameters per layer
    for (layer_name, (mins, maxs)) in layer_stats.iter() {
        let global_min = mins.iter().cloned().fold(f32::INFINITY, f32::min);
        let global_max = maxs.iter().cloned().fold(f32::NEG_INFINITY, f32::max);

        // Calculate INT8 quantization parameters (symmetric)
        let abs_max = global_min.abs().max(global_max.abs());
        let scale = abs_max / 127.0;
        let zero_point = 127i8; // Symmetric quantization

        println!("✅ Layer {}: scale={:.6}, zero_point={}", layer_name, scale, zero_point);

        // Verify valid parameters
        assert!(scale > 0.0 && scale.is_finite(), "Invalid scale for layer {}", layer_name);
        assert_eq!(zero_point, 127i8, "Symmetric quantization should use zero_point=127");
    }

    Ok(())
}

/// Test 5: Save calibration parameters to JSON
#[tokio::test]
async fn test_save_calibration_to_json() -> Result<(), Box<dyn std::error::Error>> {
    use serde::{Serialize, Deserialize};
    use std::collections::HashMap;

    #[derive(Serialize, Deserialize, Debug)]
    struct LayerQuantizationParams {
        scale: f32,
        zero_point: i8,
        min_val: f32,
        max_val: f32,
    }

    #[derive(Serialize, Deserialize, Debug)]
    struct CalibrationData {
        num_samples: usize,
        layers: HashMap<String, LayerQuantizationParams>,
    }

    // Create sample calibration data
    let mut layers = HashMap::new();
    layers.insert("vsn_layer".to_string(), LayerQuantizationParams {
        scale: 0.05,
        zero_point: 127,
        min_val: -6.35,
        max_val: 6.35,
    });
    layers.insert("lstm_layer".to_string(), LayerQuantizationParams {
        scale: 0.03,
        zero_point: 127,
        min_val: -3.81,
        max_val: 3.81,
    });
    layers.insert("attention_layer".to_string(), LayerQuantizationParams {
        scale: 0.04,
        zero_point: 127,
        min_val: -5.08,
        max_val: 5.08,
    });

    let calibration_data = CalibrationData {
        num_samples: 1000,
        layers,
    };

    // Serialize to JSON
    let json_string = serde_json::to_string_pretty(&calibration_data)?;
    println!("✅ Calibration JSON:\n{}", json_string);

    // Verify JSON structure
    assert!(json_string.contains("num_samples"));
    assert!(json_string.contains("vsn_layer"));
    assert!(json_string.contains("lstm_layer"));
    assert!(json_string.contains("attention_layer"));
    assert!(json_string.contains("scale"));
    assert!(json_string.contains("zero_point"));

    // Save to file (in real implementation)
    let output_path = PathBuf::from("ml/checkpoints/tft_int8_calibration_test.json");
    if let Some(parent) = output_path.parent() {
        std::fs::create_dir_all(parent)?;
    }
    std::fs::write(&output_path, json_string)?;
    println!("✅ Saved calibration data to: {}", output_path.display());

    // Clean up test file
    let _ = std::fs::remove_file(&output_path);

    Ok(())
}

/// Test 6: End-to-end calibration workflow
#[tokio::test]
async fn test_e2e_calibration_workflow() -> Result<(), Box<dyn std::error::Error>> {
    use serde::{Serialize, Deserialize};
    use std::collections::HashMap;

    #[derive(Serialize, Deserialize, Debug)]
    struct LayerQuantizationParams {
        scale: f32,
        zero_point: i8,
        min_val: f32,
        max_val: f32,
    }

    #[derive(Serialize, Deserialize, Debug)]
    struct CalibrationData {
        num_samples: usize,
        layers: HashMap<String, LayerQuantizationParams>,
    }

    // Step 1: Load DBN data
    let dbn_file = PathBuf::from("test_data/real/databento");

    if !dbn_file.exists() {
        println!("⚠️  Skipping test: DBN directory not found");
        return Ok(());
    }

    let mut loader = DbnSequenceLoader::with_limits(60, 256, Some(10), 10).await?;
    let (train_data, _val_data) = loader.load_sequences(&dbn_file, 0.9).await?;

    println!("✅ Step 1: Loaded {} sequences", train_data.len());

    // Step 2: Create TFT model
    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);
    let config = TFTConfig {
        input_dim: 256,
        hidden_dim: 64,
        num_heads: 4,
        num_layers: 2,
        prediction_horizon: 10,
        sequence_length: 60,
        num_quantiles: 3,
        num_static_features: 2,
        num_known_features: 3,
        num_unknown_features: 256,
        batch_size: 1,
        ..Default::default()
    };

    let mut tft = TemporalFusionTransformer::new(config)?;
    println!("✅ Step 2: Created TFT model");

    // Step 3: Run calibration forward passes
    let mut layer_stats = std::collections::HashMap::new();

    for (input, _target) in train_data.iter().take(5) {
        let batch = input.dims()[0];
        let static_features = Tensor::zeros((batch, 2), DType::F32, &device)?;
        let historical_features = input.to_dtype(DType::F32)?;
        let future_features = Tensor::zeros((batch, 10, 3), DType::F32, &device)?;

        let output = tft.forward(&static_features, &historical_features, &future_features)?;

        // Collect output layer stats
        let output_vec = output.flatten_all()?.to_vec1::<f32>()?;
        let min_val = output_vec.iter().cloned().fold(f32::INFINITY, f32::min);
        let max_val = output_vec.iter().cloned().fold(f32::NEG_INFINITY, f32::max);

        let entry = layer_stats.entry("output_layer".to_string()).or_insert((Vec::new(), Vec::new()));
        entry.0.push(min_val);
        entry.1.push(max_val);
    }

    println!("✅ Step 3: Collected activation statistics");

    // Step 4: Calculate quantization parameters
    let mut calibration_layers = HashMap::new();

    for (layer_name, (mins, maxs)) in layer_stats {
        let global_min = mins.iter().cloned().fold(f32::INFINITY, f32::min);
        let global_max = maxs.iter().cloned().fold(f32::NEG_INFINITY, f32::max);

        let abs_max = global_min.abs().max(global_max.abs());
        let scale = abs_max / 127.0;
        let zero_point = 127i8;

        calibration_layers.insert(layer_name, LayerQuantizationParams {
            scale,
            zero_point,
            min_val: global_min,
            max_val: global_max,
        });
    }

    println!("✅ Step 4: Calculated quantization parameters");

    // Step 5: Save calibration data
    let calibration_data = CalibrationData {
        num_samples: train_data.len(),
        layers: calibration_layers,
    };

    let json_string = serde_json::to_string_pretty(&calibration_data)?;
    let output_path = PathBuf::from("ml/checkpoints/tft_int8_calibration_e2e_test.json");

    if let Some(parent) = output_path.parent() {
        std::fs::create_dir_all(parent)?;
    }
    std::fs::write(&output_path, &json_string)?;

    println!("✅ Step 5: Saved calibration to {}", output_path.display());

    // Verify file exists and has content
    assert!(output_path.exists(), "Calibration file not created");
    let file_size = std::fs::metadata(&output_path)?.len();
    assert!(file_size > 100, "Calibration file too small: {} bytes", file_size);

    // Clean up test file
    let _ = std::fs::remove_file(&output_path);

    println!("✅ E2E calibration workflow complete!");

    Ok(())
}