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};

/// 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(())
}