foxhunt/ml/tests/tft_test.rs

//! Temporal Fusion Transformer (TFT) Integration Tests
//!
//! Basic tests for TFT configuration and model creation.
#![allow(unused_crate_dependencies)]

use anyhow::Result;
use ml::tft::{TFTConfig, TFTState, TemporalFusionTransformer};

mod real_data_helpers;
use real_data_helpers::{load_tft_sequences, real_data_available};

/// Test TFT configuration creation with default values
#[test]
fn test_tft_config_default() -> Result<()> {
    let config = TFTConfig::default();
    assert!(config.input_dim > 0);
    assert!(config.hidden_dim > 0);
    assert!(config.num_heads > 0);
    assert!(config.num_quantiles > 0);
    assert!(config.prediction_horizon > 0);
    assert!(config.sequence_length > 0);
    Ok(())
}

/// Test TFT configuration creation with custom values
#[test]
fn test_tft_config_custom() -> Result<()> {
    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: true,
        mixed_precision: true,
        memory_efficient: true,
        max_inference_latency_us: 50,
        target_throughput_pps: 100_000,
    };

    assert_eq!(config.input_dim, 64);
    assert_eq!(config.hidden_dim, 128);
    assert_eq!(config.num_heads, 8);
    assert_eq!(config.num_quantiles, 9);
    Ok(())
}

/// Test TFT model creation
#[test]
fn test_tft_model_creation() -> Result<()> {
    let config = TFTConfig {
        input_dim: 10,
        hidden_dim: 32,
        num_heads: 4,
        num_quantiles: 5,
        prediction_horizon: 5,
        sequence_length: 20,
        num_static_features: 2,
        num_known_features: 3,
        num_unknown_features: 5,
        ..Default::default()
    };

    let tft = TemporalFusionTransformer::new(config)?;
    assert_eq!(tft.metadata.input_dim, 10);
    assert_eq!(tft.metadata.output_dim, 5);
    assert!(!tft.is_trained);
    Ok(())
}

/// Test TFT state creation
#[test]
fn test_tft_state_creation() -> Result<()> {
    let config = TFTConfig {
        hidden_dim: 32,
        sequence_length: 20,
        num_heads: 4,
        ..Default::default()
    };

    let state = TFTState::zeros(&config)?;
    assert_eq!(state.last_update, 0);
    assert!(state.attention_cache.is_empty());
    Ok(())
}

/// Test TFT performance metrics
#[test]
fn test_tft_performance_metrics() -> Result<()> {
    let config = TFTConfig {
        input_dim: 10,
        hidden_dim: 32,
        ..Default::default()
    };

    let tft = TemporalFusionTransformer::new(config)?;
    let metrics = tft.get_metrics();

    assert!(metrics.contains_key("total_inferences"));
    assert!(metrics.contains_key("avg_latency_us"));
    assert!(metrics.contains_key("max_latency_us"));
    assert!(metrics.contains_key("throughput_pps"));

    // Initial values should be zero
    assert_eq!(metrics["total_inferences"], 0.0);
    assert_eq!(metrics["avg_latency_us"], 0.0);
    Ok(())
}

/// Test TFT training state management
#[test]
fn test_tft_training_state() -> Result<()> {
    let config = TFTConfig::default();
    let mut tft = TemporalFusionTransformer::new(config)?;

    assert!(!tft.is_trained);
    tft.is_trained = true;
    assert!(tft.is_trained);
    Ok(())
}

/// Test TFT metadata
#[test]
fn test_tft_metadata() -> Result<()> {
    let config = TFTConfig {
        input_dim: 15,
        prediction_horizon: 12,
        ..Default::default()
    };

    let tft = TemporalFusionTransformer::new(config)?;
    assert_eq!(tft.metadata.input_dim, 15);
    assert_eq!(tft.metadata.output_dim, 12);
    assert_eq!(tft.metadata.version, "1.0.0");
    assert_eq!(tft.metadata.training_samples, 0);
    assert!(tft.metadata.last_trained.is_none());
    Ok(())
}

/// Test TFT configuration validation
#[test]
fn test_tft_config_validation() -> Result<()> {
    let config = TFTConfig {
        input_dim: 20,
        hidden_dim: 64,
        num_heads: 4,
        num_layers: 2,
        prediction_horizon: 10,
        sequence_length: 50,
        num_quantiles: 7,
        num_static_features: 3,
        num_known_features: 5,
        num_unknown_features: 12,
        learning_rate: 0.001,
        batch_size: 32,
        dropout_rate: 0.1,
        l2_regularization: 0.0001,
        use_flash_attention: false,
        mixed_precision: false,
        memory_efficient: true,
        max_inference_latency_us: 100,
        target_throughput_pps: 50_000,
    };

    assert!(config.input_dim > 0);
    assert!(config.hidden_dim > 0);
    assert!(config.num_heads > 0);
    assert!(config.num_layers > 0);
    assert!(config.prediction_horizon > 0);
    assert!(config.sequence_length > 0);
    assert!(config.num_quantiles > 0);
    assert!(config.learning_rate > 0.0);
    assert!(config.batch_size > 0);
    assert!(config.dropout_rate >= 0.0 && config.dropout_rate < 1.0);
    assert!(config.max_inference_latency_us > 0);
    assert!(config.target_throughput_pps > 0);
    Ok(())
}

/// Test multiple TFT model configurations
#[test]
fn test_multiple_tft_configs() -> Result<()> {
    let configs = vec![
        TFTConfig {
            input_dim: 10,
            hidden_dim: 32,
            num_heads: 2,
            ..Default::default()
        },
        TFTConfig {
            input_dim: 20,
            hidden_dim: 64,
            num_heads: 4,
            ..Default::default()
        },
        TFTConfig {
            input_dim: 30,
            hidden_dim: 128,
            num_heads: 8,
            ..Default::default()
        },
    ];

    for config in configs {
        let tft = TemporalFusionTransformer::new(config)?;
        assert!(!tft.is_trained);
        assert!(tft.metadata.training_samples == 0);
    }
    Ok(())
}

// ============================================================================
// Real Market Data Tests
// ============================================================================

/// Test TFT model creation with real market data dimensions
#[tokio::test]
async fn test_tft_model_creation_real_data_dimensions() -> Result<()> {
    // Skip if no real data available
    if !real_data_available().await {
        eprintln!("Skipping test: real data not available");
        return Ok(());
    }

    // Load sample sequences to determine realistic dimensions
    let sequences = load_tft_sequences(10, 20, 10).await?;
    if sequences.is_empty() {
        eprintln!("Skipping test: no sequences loaded");
        return Ok(());
    }

    // Create TFT config matching real data dimensions
    let config = TFTConfig {
        input_dim: 10,            // Features per timestep (OHLCV)
        hidden_dim: 64,
        num_heads: 4,
        num_quantiles: 5,
        prediction_horizon: 5,
        sequence_length: 20,      // 20 timesteps history
        num_static_features: 2,
        num_known_features: 3,
        num_unknown_features: 5,
        ..Default::default()
    };

    let tft = TemporalFusionTransformer::new(config)?;
    assert_eq!(tft.metadata.input_dim, 10);
    assert_eq!(tft.metadata.output_dim, 5);
    assert!(!tft.is_trained);

    Ok(())
}

/// Test TFT state creation with real market data sequence length
#[tokio::test]
async fn test_tft_state_creation_real_data() -> Result<()> {
    // Skip if no real data available
    if !real_data_available().await {
        eprintln!("Skipping test: real data not available");
        return Ok(());
    }

    let config = TFTConfig {
        hidden_dim: 64,
        sequence_length: 20,  // Real data sequence length
        num_heads: 4,
        ..Default::default()
    };

    let state = TFTState::zeros(&config)?;
    assert_eq!(state.last_update, 0);
    assert!(state.attention_cache.is_empty());

    Ok(())
}

/// Test TFT configuration validation with real data parameters
#[tokio::test]
async fn test_tft_config_validation_real_data() -> Result<()> {
    // Skip if no real data available
    if !real_data_available().await {
        eprintln!("Skipping test: real data not available");
        return Ok(());
    }

    // Load sequences to verify configuration matches data
    let sequences = load_tft_sequences(30, 50, 10).await?;
    if sequences.is_empty() {
        eprintln!("Skipping test: no sequences loaded");
        return Ok(());
    }

    let config = TFTConfig {
        input_dim: 10,             // Match real data features
        hidden_dim: 64,
        num_heads: 4,
        num_layers: 2,
        prediction_horizon: 10,
        sequence_length: 50,       // Match loaded sequence length
        num_quantiles: 7,
        num_static_features: 3,
        num_known_features: 5,
        num_unknown_features: 12,
        learning_rate: 0.001,
        batch_size: 32,
        dropout_rate: 0.1,
        l2_regularization: 0.0001,
        use_flash_attention: false,
        mixed_precision: false,
        memory_efficient: true,
        max_inference_latency_us: 100,
        target_throughput_pps: 50_000,
    };

    // Validate all config parameters
    assert!(config.input_dim > 0);
    assert!(config.hidden_dim > 0);
    assert!(config.num_heads > 0);
    assert!(config.num_layers > 0);
    assert!(config.prediction_horizon > 0);
    assert!(config.sequence_length > 0);
    assert!(config.num_quantiles > 0);
    assert!(config.learning_rate > 0.0);
    assert!(config.batch_size > 0);
    assert!(config.dropout_rate >= 0.0 && config.dropout_rate < 1.0);
    assert!(config.max_inference_latency_us > 0);
    assert!(config.target_throughput_pps > 0);

    // Create model with validated config
    let tft = TemporalFusionTransformer::new(config)?;
    assert_eq!(tft.metadata.input_dim, 10);
    assert_eq!(tft.metadata.output_dim, 10);
    assert!(!tft.is_trained);

    Ok(())
}