foxhunt/ml/tests/tft_int8_training_pipeline_test.rs

//! TFT INT8 Training Pipeline Tests (TDD)
//!
//! Tests for training TFT model with ES.FUT data and applying INT8 quantization
//! using calibration data from Agent 10.3.
//!
//! **TDD Phases**:
//! 1. RED: Write test first (should FAIL)
//! 2. GREEN: Implement to make test PASS
//! 3. REFACTOR: Add comprehensive tests (7+ total)

use anyhow::Result;
use candle_core::{Device, DType};
use std::sync::Arc;

use ml::checkpoint::FileSystemStorage;
use ml::memory_optimization::quantization::{
    extract_weights_from_varmap, Quantizer, QuantizationConfig, QuantizationType,
};
use ml::tft::{TFTConfig, TemporalFusionTransformer};
use ml::trainers::tft::{TFTTrainer, TFTTrainerConfig};
use ml::tft::training::TFTDataLoader;
use ml::ModelType;

// Import DBN loading utilities from train_tft_dbn.rs example
use dbn::decode::{DecodeRecordRef, DbnDecoder};
use dbn::OhlcvMsg;
use chrono::{DateTime, TimeZone, Utc};
use ndarray::{Array1, Array2};

/// OHLCV bar structure (intermediate format)
#[derive(Debug, Clone)]
struct OhlcvBar {
    timestamp: DateTime<Utc>,
    open: f64,
    high: f64,
    low: f64,
    close: f64,
    volume: f64,
}

/// Load OHLCV bars from DBN file (simplified version from train_tft_dbn.rs)
async fn load_dbn_ohlcv_bars(file_path: &str) -> Result<Vec<OhlcvBar>> {
    let mut decoder = DbnDecoder::from_file(file_path)?;
    let mut bars = Vec::new();
    let mut prev_close: Option<f64> = None;

    while let Some(record_ref) = decoder.decode_record_ref()? {
        if let Some(ohlcv) = record_ref.get::<OhlcvMsg>() {
            let ts_nanos = ohlcv.hd.ts_event as i64;
            let secs = ts_nanos / 1_000_000_000;
            let nanos = (ts_nanos % 1_000_000_000) as u32;
            let timestamp = Utc.timestamp_opt(secs, nanos).single()
                .ok_or_else(|| anyhow::anyhow!("Invalid timestamp"))?;

            let mut open_f64 = ohlcv.open as f64 / 1_000_000_000.0;
            let mut high_f64 = ohlcv.high as f64 / 1_000_000_000.0;
            let mut low_f64 = ohlcv.low as f64 / 1_000_000_000.0;
            let mut close_f64 = ohlcv.close as f64 / 1_000_000_000.0;

            // Price anomaly correction
            if let Some(prev) = prev_close {
                let pct_change = ((close_f64 - prev) / prev).abs();
                if pct_change > 0.5 && close_f64 < 1000.0 {
                    let corrected_close = close_f64 * 100.0;
                    if corrected_close >= 3000.0 && corrected_close <= 6000.0 {
                        open_f64 *= 100.0;
                        high_f64 *= 100.0;
                        low_f64 *= 100.0;
                        close_f64 = corrected_close;
                    } else {
                        prev_close = Some(prev);
                        continue;
                    }
                }
            }

            prev_close = Some(close_f64);

            bars.push(OhlcvBar {
                timestamp,
                open: open_f64,
                high: high_f64,
                low: low_f64,
                close: close_f64,
                volume: ohlcv.volume as f64,
            });
        }
    }

    Ok(bars)
}

/// Convert OHLCV bars to TFT format (simplified version)
fn convert_to_tft_data(
    bars: &[OhlcvBar],
    lookback: usize,
    horizon: usize,
) -> Result<Vec<(Array1<f64>, Array2<f64>, Array2<f64>, Array1<f64>)>> {
    if bars.len() < lookback + horizon {
        anyhow::bail!("Not enough data");
    }

    let mut samples = Vec::new();
    let mean_price = bars.iter().map(|b| b.close).sum::<f64>() / bars.len() as f64;
    let mean_volume = bars.iter().map(|b| b.volume).sum::<f64>() / bars.len() as f64;

    for i in 0..bars.len() - lookback - horizon + 1 {
        // Static features (10)
        let static_feat = Array1::from_vec(vec![
            mean_price / 5000.0, 0.01, mean_volume / 1000.0, 0.01,
            0.5, 0.5, 0.5, 0.5, 0.01, 0.01,
        ]);

        // Historical features (lookback x 50)
        let mut hist_data = Vec::new();
        for t in 0..lookback {
            let bar = &bars[i + t];
            let mut features = vec![
                bar.open / mean_price, bar.high / mean_price, bar.low / mean_price,
                bar.close / mean_price, bar.volume / mean_volume,
            ];
            // Pad to 50 features
            features.extend(vec![0.0; 45]);
            hist_data.extend(features);
        }
        let hist_feat = Array2::from_shape_vec((lookback, 50), hist_data)?;

        // Future features (horizon x 10)
        let fut_data = vec![0.5; horizon * 10];
        let fut_feat = Array2::from_shape_vec((horizon, 10), fut_data)?;

        // Targets (horizon)
        let targets: Vec<f64> = (0..horizon)
            .map(|t| bars[i + lookback + t].close / mean_price)
            .collect();
        let target_arr = Array1::from_vec(targets);

        samples.push((static_feat, hist_feat, fut_feat, target_arr));
    }

    Ok(samples)
}

// ============================================================================
// TDD Phase 1: RED - Write failing test
// ============================================================================

/// Test 1: Train TFT and apply INT8 quantization (PRIMARY TEST - SHOULD FAIL)
#[tokio::test]
#[ignore] // Remove this after implementation
async fn test_tft_trains_and_quantizes() -> Result<()> {
    // Use absolute path from project root (one level up from ml/)
    let ml_dir = std::env::current_dir()?;
    let project_root = ml_dir.parent().unwrap_or(&ml_dir);
    let dbn_file = project_root.join("test_data/real/databento/ES.FUT_ohlcv-1m_2024-01-02.dbn");

    if !dbn_file.exists() {
        eprintln!("⚠️  Skipping test: DBN file not found at {:?}", dbn_file);
        eprintln!("    Current dir: {:?}", project_root);
        return Ok(());
    }

    let dbn_file_str = dbn_file.to_str().unwrap();

    // Load ES.FUT data
    println!("📊 Loading ES.FUT data from: {:?}", dbn_file);
    let bars = load_dbn_ohlcv_bars(dbn_file_str).await?;
    println!("✅ Loaded {} bars", bars.len());

    // Convert to TFT format
    let lookback = 26; // Reduced from 60 for faster testing
    let horizon = 10;
    let tft_data = convert_to_tft_data(&bars, lookback, horizon)?;
    println!("✅ Created {} TFT samples", tft_data.len());

    // Split train/val
    let split_idx = (tft_data.len() as f64 * 0.8) as usize;
    let train_data = tft_data[..split_idx].to_vec();
    let val_data = tft_data[split_idx..].to_vec();

    // Train TFT model (F32)
    println!("\n🏋️  Training TFT model (F32) for 10 epochs...");
    let trainer_config = TFTTrainerConfig {
        epochs: 10,
        learning_rate: 0.001,
        batch_size: 16, // Reduced for faster testing
        hidden_dim: 128,
        num_attention_heads: 4,
        dropout_rate: 0.1,
        lstm_layers: 2,
        quantiles: vec![0.1, 0.5, 0.9],
        lookback_window: lookback,
        forecast_horizon: horizon,
        use_gpu: Device::cuda_if_available(0).is_ok(),
        checkpoint_dir: "ml/checkpoints/tft_test".to_string(),
    };

    let storage = Arc::new(FileSystemStorage::new(
        std::path::PathBuf::from(&trainer_config.checkpoint_dir)
    ));
    let mut trainer = TFTTrainer::new(trainer_config.clone(), storage)?;

    let train_loader = TFTDataLoader::new(train_data, trainer_config.batch_size, true);
    let val_loader = TFTDataLoader::new(val_data, trainer_config.batch_size, false);

    let final_metrics = trainer.train(train_loader, val_loader).await?;
    println!("✅ Training complete - Val Loss: {:.6}", final_metrics.val_loss);

    // Load calibration data (Agent 10.3)
    println!("\n📊 Loading calibration data...");
    let calibration_path = project_root.join("ml/calibration/es_fut_calibration.json");
    if !calibration_path.exists() {
        anyhow::bail!("❌ Calibration file not found: {:?}", calibration_path);
    }

    let calibration_json = std::fs::read_to_string(&calibration_path)?;
    let calibration: serde_json::Value = serde_json::from_str(&calibration_json)?;
    let sample_count = calibration["samples"].as_array()
        .ok_or_else(|| anyhow::anyhow!("Invalid calibration format"))?
        .len();
    println!("✅ Loaded {} calibration samples", sample_count);

    // Apply INT8 quantization using VarMap extraction
    println!("\n🔧 Applying INT8 quantization...");
    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);
    
    // Get trained model's VarMap
    let model = trainer.get_model();
    let varmap = model.get_varmap();

    // Extract key weights and quantize
    let config = QuantizationConfig {
        quant_type: QuantizationType::Int8,
        symmetric: true,
        per_channel: false,
        calibration_samples: Some(sample_count),
    };
    let mut quantizer = Quantizer::new(config, device);

    // Extract and quantize attention weights (example)
    let attention_weight = extract_weights_from_varmap(
        &varmap,
        "temporal_attention.query_proj.weight"
    )?;
    let quantized_attn = quantizer.quantize_tensor(&attention_weight, "attn.weight")?;

    println!("✅ Quantization complete:");
    println!("   • Type: {:?}", quantized_attn.quant_type);
    println!("   • Scale: {:.6}", quantized_attn.scale);
    println!("   • Memory savings: {:.2} MB", quantizer.memory_savings_mb());

    // Verify accuracy loss <10% (relaxed for test)
    let dequantized = quantizer.dequantize_tensor(&quantized_attn)?;
    let original_norm = attention_weight.sqr()?.sum_all()?.to_vec0::<f32>()?;
    let dequant_norm = dequantized.sqr()?.sum_all()?.to_vec0::<f32>()?;
    let accuracy_loss = ((original_norm - dequant_norm).abs() / original_norm) * 100.0;

    println!("\n📊 Accuracy Metrics:");
    println!("   • Accuracy loss: {:.2}%", accuracy_loss);
    assert!(accuracy_loss < 10.0, "Accuracy loss too high: {:.2}%", accuracy_loss);

    println!("\n✅ TFT INT8 training pipeline test PASSED");
    Ok(())
}

// ============================================================================
// TDD Phase 3: REFACTOR - Add comprehensive tests
// ============================================================================

/// Test 2: F32 training only (baseline)
#[tokio::test]
#[ignore]
async fn test_tft_f32_training_only() -> Result<()> {
    // Test F32 training without quantization
    Ok(())
}

/// Test 3: Quantization accuracy (isolated)
#[tokio::test]
#[ignore]
async fn test_int8_quantization_accuracy() -> Result<()> {
    // Test quantization accuracy in isolation
    Ok(())
}

/// Test 4: INT8 inference (dequantize on-the-fly)
#[tokio::test]
#[ignore]
async fn test_int8_inference() -> Result<()> {
    // Test INT8 inference with dequantization
    Ok(())
}

/// Test 5: Memory reduction (75%)
#[tokio::test]
#[ignore]
async fn test_memory_reduction() -> Result<()> {
    // Verify 75% memory reduction
    Ok(())
}

/// Test 6: Checkpoint persistence
#[tokio::test]
#[ignore]
async fn test_checkpoint_save_load() -> Result<()> {
    // Test saving and loading both F32 and INT8 checkpoints
    Ok(())
}

/// Test 7: Calibration data integration
#[tokio::test]
#[ignore]
async fn test_calibration_integration() -> Result<()> {
    // Verify calibration data is properly used
    Ok(())
}

/// Test 8: End-to-end pipeline (full integration)
#[tokio::test]
#[ignore]
async fn test_e2e_training_quantization_inference() -> Result<()> {
    // Full pipeline: train → quantize → save → load → infer
    Ok(())
}