foxhunt/ml/tests/tft_test.rs

use crate::MLError;
use candle_core::{DType, Device, Tensor};
use ml::tft::attention::{AttentionWeights, MultiHeadAttention, TemporalAttention};
use ml::tft::variable_selection::{
    VariableSelection, VariableSelectionOutput, VariableSelectionWeights,
};
use ml::tft::{QuantileOutput, TFTConfig, TFTModel, TemporalFusionTransformer};
use proptest::prelude::*;
use std::collections::HashMap;
use tokio;
use common::{ModelPerformance, TimeSeriesData, TradingSignal};

/// Mock TFT Model for testing
#[derive(Debug)]
pub struct MockTFTModel {
    pub config: TFTConfig,
    pub variable_selection: MockVariableSelectionNetwork,
    pub temporal_attention: MockTemporalAttention,
    pub forward_calls: usize,
    pub training_steps: usize,
    pub quantile_predictions: Vec<QuantileOutput>,
}

impl MockTFTModel {
    pub fn new(config: TFTConfig) -> Self {
        Self {
            config: config.clone(),
            variable_selection: MockVariableSelectionNetwork::new(&config),
            temporal_attention: MockTemporalAttention::new(&config),
            forward_calls: 0,
            training_steps: 0,
            quantile_predictions: Vec::new(),
        }
    }

    pub async fn forward(&mut self, input: &TimeSeriesData) -> Result<QuantileOutput, MLError> {
        self.forward_calls += 1;

        // Variable selection step
        let selected_features = self
            .variable_selection
            .select_variables(&input.features)
            .await?;

        // Temporal attention step
        let attention_weights = self
            .temporal_attention
            .compute_attention(&selected_features, input.sequence_length)
            .await?;

        // Generate quantile predictions
        let quantile_output = self
            .generate_quantile_predictions(&selected_features, &attention_weights)
            .await?;

        self.quantile_predictions.push(quantile_output.clone());
        Ok(quantile_output)
    }

    async fn generate_quantile_predictions(
        &self,
        features: &[f32],
        attention_weights: &[f32],
    ) -> Result<QuantileOutput, MLError> {
        if features.is_empty() || attention_weights.is_empty() {
            return Err(MLError::InvalidInput(
                "Empty features or attention weights".to_string(),
            ));
        }

        let mut predictions = HashMap::new();

        // Generate predictions for different quantiles
        for &quantile in &self.config.quantiles {
            let mut prediction = 0.0;

            // Weighted combination of features
            for (i, &feature) in features.iter().enumerate() {
                let weight = if i < attention_weights.len() {
                    attention_weights[i]
                } else {
                    1.0
                };
                prediction += feature * weight * (quantile as f32); // Mock quantile-specific prediction
            }

            // Add some quantile-specific adjustment
            prediction *= if quantile < 0.5 { 0.9 } else { 1.1 };
            predictions.insert((quantile * 100.0) as u8, prediction);
        }

        Ok(QuantileOutput {
            predictions,
            prediction_intervals: self.compute_prediction_intervals(&predictions),
            point_forecast: predictions.get(&50).cloned().unwrap_or(0.0), // Median as point forecast
            uncertainty_score: self.compute_uncertainty(&predictions),
        })
    }

    fn compute_prediction_intervals(
        &self,
        predictions: &HashMap<u8, f32>,
    ) -> HashMap<String, (f32, f32)> {
        let mut intervals = HashMap::new();

        // Common prediction intervals
        if let (Some(&p10), Some(&p90)) = (predictions.get(&10), predictions.get(&90)) {
            intervals.insert("80%".to_string(), (p10, p90));
        }
        if let (Some(&p5), Some(&p95)) = (predictions.get(&5), predictions.get(&95)) {
            intervals.insert("90%".to_string(), (p5, p95));
        }
        if let (Some(&p25), Some(&p75)) = (predictions.get(&25), predictions.get(&75)) {
            intervals.insert("50%".to_string(), (p25, p75));
        }

        intervals
    }

    fn compute_uncertainty(&self, predictions: &HashMap<u8, f32>) -> f32 {
        if let (Some(&p10), Some(&p90)) = (predictions.get(&10), predictions.get(&90)) {
            (p90 - p10).abs() / 2.0 // Width of 80% prediction interval
        } else {
            0.0
        }
    }

    pub async fn train(
        &mut self,
        batch: &[TimeSeriesData],
        targets: &[Vec<f32>],
    ) -> Result<f32, MLError> {
        self.training_steps += 1;

        let mut total_loss = 0.0;
        let batch_size = batch.len();

        for (input, target) in batch.iter().zip(targets.iter()) {
            let prediction = self.forward(input).await?;

            // Compute quantile loss for each quantile
            for &quantile in &self.config.quantiles {
                let quantile_key = (quantile * 100.0) as u8;
                if let Some(&pred) = prediction.predictions.get(&quantile_key) {
                    for &actual in target {
                        let error = actual - pred;
                        let quantile_loss = if error >= 0.0 {
                            quantile * error
                        } else {
                            (quantile - 1.0) * error
                        };
                        total_loss += quantile_loss.abs();
                    }
                }
            }
        }

        // Return decreasing loss over training steps
        let avg_loss = total_loss / (batch_size as f32 * self.config.quantiles.len() as f32);
        Ok(avg_loss / (self.training_steps as f32 + 1.0))
    }

    pub async fn multi_horizon_predict(
        &mut self,
        input: &TimeSeriesData,
        horizons: &[usize],
    ) -> Result<HashMap<usize, QuantileOutput>, MLError> {
        let mut predictions = HashMap::new();

        for &horizon in horizons {
            // Modify input for specific horizon prediction
            let mut horizon_input = input.clone();
            horizon_input.prediction_horizon = horizon;

            let prediction = self.forward(&horizon_input).await?;
            predictions.insert(horizon, prediction);
        }

        Ok(predictions)
    }

    pub fn get_variable_importance(&self) -> HashMap<usize, f64> {
        self.variable_selection.get_importance_scores()
    }

    pub fn get_attention_weights(&self) -> Vec<f32> {
        self.temporal_attention.get_latest_weights()
    }
}

/// Mock Variable Selection Network
#[derive(Debug)]
pub struct MockVariableSelectionNetwork {
    pub importance_scores: HashMap<usize, f64>,
    pub selection_threshold: f64,
    pub num_variables: usize,
}

impl MockVariableSelectionNetwork {
    pub fn new(config: &TFTConfig) -> Self {
        let mut importance_scores = HashMap::new();

        // Initialize with random importance scores
        for i in 0..config.num_input_features {
            let importance = (i as f64 + 1.0) / (config.num_input_features as f64 + 1.0); // Decreasing importance
            importance_scores.insert(i, importance);
        }

        Self {
            importance_scores,
            selection_threshold: 0.1, // Only select features with importance > 0.1
            num_variables: config.num_input_features,
        }
    }

    pub async fn select_variables(&mut self, features: &[f32]) -> Result<Vec<f32>, MLError> {
        if features.len() != self.num_variables {
            return Err(MLError::DimensionMismatch(format!(
                "Expected {} features, got {}",
                self.num_variables,
                features.len()
            )));
        }

        let mut selected_features = Vec::new();

        for (i, &feature) in features.iter().enumerate() {
            if let Some(&importance) = self.importance_scores.get(&i) {
                if importance > self.selection_threshold {
                    // Weight feature by its importance
                    selected_features.push(feature * importance as f32);
                }
            }
        }

        if selected_features.is_empty() {
            // If no features selected, use top feature
            selected_features.push(features[0]);
        }

        Ok(selected_features)
    }

    pub fn get_importance_scores(&self) -> HashMap<usize, f64> {
        self.importance_scores.clone()
    }

    pub fn update_importance_scores(&mut self, new_scores: HashMap<usize, f64>) {
        self.importance_scores = new_scores;
    }
}

/// Mock Temporal Attention mechanism
#[derive(Debug)]
pub struct MockTemporalAttention {
    pub num_heads: usize,
    pub attention_dim: usize,
    pub latest_weights: Vec<f32>,
}

impl MockTemporalAttention {
    pub fn new(config: &TFTConfig) -> Self {
        Self {
            num_heads: config.num_attention_heads,
            attention_dim: config.attention_dim,
            latest_weights: Vec::new(),
        }
    }

    pub async fn compute_attention(
        &mut self,
        features: &[f32],
        sequence_length: usize,
    ) -> Result<Vec<f32>, MLError> {
        if features.is_empty() {
            return Err(MLError::InvalidInput(
                "Empty features for attention computation".to_string(),
            ));
        }

        let effective_seq_len = sequence_length.min(features.len());
        let mut attention_weights = Vec::with_capacity(effective_seq_len);

        // Mock attention computation - in practice, this would be multi-head self-attention
        for i in 0..effective_seq_len {
            // Simple attention: recent positions get higher weights
            let position_weight = (i as f32 + 1.0) / effective_seq_len as f32;

            // Feature-dependent attention
            let feature_magnitude = if i < features.len() {
                features[i].abs()
            } else {
                1.0
            };

            // Combined attention score
            let attention_score = position_weight * (1.0 + feature_magnitude);
            attention_weights.push(attention_score);
        }

        // Normalize attention weights to sum to 1
        let sum: f32 = attention_weights.iter().sum();
        if sum > 0.0 {
            for weight in &mut attention_weights {
                *weight /= sum;
            }
        }

        self.latest_weights = attention_weights.clone();
        Ok(attention_weights)
    }

    pub fn get_latest_weights(&self) -> Vec<f32> {
        self.latest_weights.clone()
    }

    pub async fn multi_head_attention(
        &mut self,
        features: &[f32],
        sequence_length: usize,
    ) -> Result<Vec<Vec<f32>>, MLError> {
        let mut head_weights = Vec::new();

        for head in 0..self.num_heads {
            // Each head focuses on different aspects
            let head_features: Vec<f32> = features
                .iter()
                .enumerate()
                .map(|(i, &f)| {
                    f * ((head + 1) as f32 / self.num_heads as f32) + (i % (head + 1)) as f32 * 0.1
                })
                .collect();

            let weights = self
                .compute_attention(&head_features, sequence_length)
                .await?;
            head_weights.push(weights);
        }

        Ok(head_weights)
    }
}

/// Mock time series data for testing
pub fn create_mock_time_series(length: usize, num_features: usize) -> TimeSeriesData {
    let mut features = Vec::new();

    for i in 0..length {
        let mut row = Vec::new();
        for j in 0..num_features {
            // Generate synthetic time series with trend and seasonality
            let trend = (i as f32) * 0.01;
            let seasonality = (i as f32 * 2.0 * std::f32::consts::PI / 24.0).sin() * 0.5;
            let noise = (i as f32 * j as f32).sin() * 0.1;
            row.push(trend + seasonality + noise);
        }
        features.push(row);
    }

    TimeSeriesData {
        features,
        sequence_length: length,
        prediction_horizon: 10,
        target_column: 0,
        timestamp: std::time::SystemTime::now(),
    }
}

#[tokio::test]
async fn test_tft_model_creation() {
    let config = TFTConfig {
        num_input_features: 20,
        hidden_dim: 256,
        num_attention_heads: 8,
        attention_dim: 64,
        num_quantiles: 9,
        quantiles: vec![0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9],
        dropout: 0.1,
        max_sequence_length: 168,                // 1 week of hourly data
        prediction_horizons: vec![1, 6, 12, 24], // 1h, 6h, 12h, 24h ahead
        use_static_features: true,
        use_temporal_features: true,
        variable_selection_threshold: 0.1,
        attention_temperature: 1.0,
    };

    let model = MockTFTModel::new(config.clone());
    assert_eq!(model.config.num_input_features, 20);
    assert_eq!(model.config.hidden_dim, 256);
    assert_eq!(model.config.num_attention_heads, 8);
    assert_eq!(model.config.quantiles.len(), 9);
    assert_eq!(model.forward_calls, 0);
    assert_eq!(model.training_steps, 0);
}

#[tokio::test]
async fn test_variable_selection() {
    let config = TFTConfig {
        num_input_features: 10,
        hidden_dim: 128,
        num_attention_heads: 4,
        attention_dim: 32,
        num_quantiles: 5,
        quantiles: vec![0.1, 0.3, 0.5, 0.7, 0.9],
        dropout: 0.1,
        max_sequence_length: 100,
        prediction_horizons: vec![1, 5, 10],
        use_static_features: true,
        use_temporal_features: true,
        variable_selection_threshold: 0.3, // Higher threshold
        attention_temperature: 1.0,
    };

    let mut vsn = MockVariableSelectionNetwork::new(&config);
    let features = vec![1.0, 0.5, -0.3, 2.0, 0.1, -1.5, 0.8, -0.2, 1.2, 0.9];

    let selected = vsn.select_variables(&features).await.unwrap();

    // Should select features with importance > 0.3
    assert!(!selected.is_empty());
    assert!(selected.len() <= features.len()); // Should select subset

    let importance_scores = vsn.get_importance_scores();
    assert_eq!(importance_scores.len(), 10);

    // Importance scores should be in [0, 1] range
    for &importance in importance_scores.values() {
        assert!(importance >= 0.0 && importance <= 1.0);
    }
}

#[tokio::test]
async fn test_temporal_attention() {
    let config = TFTConfig {
        num_input_features: 8,
        hidden_dim: 64,
        num_attention_heads: 2,
        attention_dim: 32,
        num_quantiles: 3,
        quantiles: vec![0.25, 0.5, 0.75],
        dropout: 0.05,
        max_sequence_length: 50,
        prediction_horizons: vec![1, 5],
        use_static_features: false,
        use_temporal_features: true,
        variable_selection_threshold: 0.1,
        attention_temperature: 1.0,
    };

    let mut attention = MockTemporalAttention::new(&config);
    let features = vec![1.0, 0.8, 0.6, 0.4, 0.2, 0.9, 0.7, 0.5];
    let sequence_length = 8;

    let weights = attention
        .compute_attention(&features, sequence_length)
        .await
        .unwrap();

    assert_eq!(weights.len(), sequence_length);

    // Attention weights should sum to approximately 1.0
    let sum: f32 = weights.iter().sum();
    assert!((sum - 1.0).abs() < 1e-6);

    // All weights should be non-negative
    for &weight in &weights {
        assert!(weight >= 0.0);
        assert!(weight.is_finite());
    }
}

#[tokio::test]
async fn test_multi_head_attention() {
    let config = TFTConfig {
        num_input_features: 6,
        hidden_dim: 48,
        num_attention_heads: 3, // Multiple heads
        attention_dim: 16,
        num_quantiles: 5,
        quantiles: vec![0.1, 0.25, 0.5, 0.75, 0.9],
        dropout: 0.1,
        max_sequence_length: 24,
        prediction_horizons: vec![1, 3, 6],
        use_static_features: true,
        use_temporal_features: true,
        variable_selection_threshold: 0.2,
        attention_temperature: 1.0,
    };

    let mut attention = MockTemporalAttention::new(&config);
    let features = vec![1.5, -0.5, 2.0, 0.3, -1.0, 0.8];
    let sequence_length = 6;

    let head_weights = attention
        .multi_head_attention(&features, sequence_length)
        .await
        .unwrap();

    assert_eq!(head_weights.len(), 3); // 3 attention heads

    for head_weight in head_weights {
        assert_eq!(head_weight.len(), sequence_length);

        // Each head's weights should sum to 1.0
        let sum: f32 = head_weight.iter().sum();
        assert!((sum - 1.0).abs() < 1e-5);

        // All weights should be valid
        for &weight in &head_weight {
            assert!(weight >= 0.0);
            assert!(weight.is_finite());
        }
    }
}

#[tokio::test]
async fn test_quantile_prediction() {
    let config = TFTConfig {
        num_input_features: 5,
        hidden_dim: 32,
        num_attention_heads: 2,
        attention_dim: 16,
        num_quantiles: 7,
        quantiles: vec![0.05, 0.1, 0.25, 0.5, 0.75, 0.9, 0.95],
        dropout: 0.1,
        max_sequence_length: 20,
        prediction_horizons: vec![1, 5, 10],
        use_static_features: true,
        use_temporal_features: true,
        variable_selection_threshold: 0.1,
        attention_temperature: 1.0,
    };

    let mut model = MockTFTModel::new(config);
    let time_series = create_mock_time_series(15, 5);

    let prediction = model.forward(&time_series).await.unwrap();

    // Check quantile predictions
    assert_eq!(prediction.predictions.len(), 7);
    assert!(prediction.predictions.contains_key(&5)); // 0.05 quantile
    assert!(prediction.predictions.contains_key(&50)); // 0.5 quantile (median)
    assert!(prediction.predictions.contains_key(&95)); // 0.95 quantile

    // Check prediction intervals
    assert!(prediction.prediction_intervals.len() > 0);
    if let Some(&(lower, upper)) = prediction.prediction_intervals.get("90%") {
        assert!(lower <= upper); // Lower bound should be <= upper bound
    }

    // Point forecast should be the median
    assert!((prediction.point_forecast - prediction.predictions.get(&50).unwrap()).abs() < 1e-6);

    // Uncertainty score should be non-negative
    assert!(prediction.uncertainty_score >= 0.0);
    assert!(prediction.uncertainty_score.is_finite());
}

#[tokio::test]
async fn test_multi_horizon_prediction() {
    let config = TFTConfig {
        num_input_features: 4,
        hidden_dim: 24,
        num_attention_heads: 2,
        attention_dim: 12,
        num_quantiles: 5,
        quantiles: vec![0.1, 0.25, 0.5, 0.75, 0.9],
        dropout: 0.05,
        max_sequence_length: 12,
        prediction_horizons: vec![1, 3, 6, 12],
        use_static_features: false,
        use_temporal_features: true,
        variable_selection_threshold: 0.15,
        attention_temperature: 0.8,
    };

    let mut model = MockTFTModel::new(config);
    let time_series = create_mock_time_series(10, 4);
    let horizons = vec![1, 3, 6];

    let predictions = model
        .multi_horizon_predict(&time_series, &horizons)
        .await
        .unwrap();

    assert_eq!(predictions.len(), 3);
    assert!(predictions.contains_key(&1));
    assert!(predictions.contains_key(&3));
    assert!(predictions.contains_key(&6));

    // Each horizon should have valid predictions
    for (horizon, prediction) in predictions {
        assert_eq!(prediction.predictions.len(), 5); // 5 quantiles
        assert!(prediction.point_forecast.is_finite());
        assert!(prediction.uncertainty_score >= 0.0);

        // Longer horizons might have higher uncertainty (not guaranteed, but often true)
        if horizon > 1 {
            // Just check that uncertainty is reasonable
            assert!(prediction.uncertainty_score < 100.0); // Reasonable bound
        }
    }
}

#[tokio::test]
async fn test_tft_training() {
    let config = TFTConfig {
        num_input_features: 3,
        hidden_dim: 16,
        num_attention_heads: 2,
        attention_dim: 8,
        num_quantiles: 3,
        quantiles: vec![0.25, 0.5, 0.75],
        dropout: 0.1,
        max_sequence_length: 8,
        prediction_horizons: vec![1, 2],
        use_static_features: false,
        use_temporal_features: true,
        variable_selection_threshold: 0.1,
        attention_temperature: 1.0,
    };

    let mut model = MockTFTModel::new(config);

    // Create training batch
    let batch = vec![
        create_mock_time_series(6, 3),
        create_mock_time_series(6, 3),
        create_mock_time_series(6, 3),
    ];

    let targets = vec![
        vec![1.0, 1.1, 1.2], // Target values for first sample
        vec![0.8, 0.9, 1.0], // Target values for second sample
        vec![1.2, 1.3, 1.4], // Target values for third sample
    ];

    // Perform training steps
    let mut losses = Vec::new();
    for _ in 0..5 {
        let loss = model.train(&batch, &targets).await.unwrap();
        losses.push(loss);
    }

    assert_eq!(model.training_steps, 5);
    assert!(losses[0] > 0.0);
    assert!(losses[4] < losses[0]); // Loss should decrease over training

    // Should have made predictions during training
    assert!(!model.quantile_predictions.is_empty());
}

#[tokio::test]
async fn test_variable_importance_analysis() {
    let config = TFTConfig {
        num_input_features: 8,
        hidden_dim: 32,
        num_attention_heads: 2,
        attention_dim: 16,
        num_quantiles: 3,
        quantiles: vec![0.3, 0.5, 0.7],
        dropout: 0.1,
        max_sequence_length: 16,
        prediction_horizons: vec![1, 4],
        use_static_features: true,
        use_temporal_features: true,
        variable_selection_threshold: 0.2,
        attention_temperature: 1.0,
    };

    let mut model = MockTFTModel::new(config);
    let time_series = create_mock_time_series(12, 8);

    // Make prediction to update importance scores
    let _ = model.forward(&time_series).await.unwrap();

    let importance_scores = model.get_variable_importance();
    assert_eq!(importance_scores.len(), 8);

    // Check that importance scores are reasonable
    let max_importance = importance_scores.values().fold(0.0, |acc, &x| acc.max(x));
    let min_importance = importance_scores.values().fold(1.0, |acc, &x| acc.min(x));

    assert!(max_importance > min_importance); // Should have variation
    assert!(max_importance <= 1.0);
    assert!(min_importance >= 0.0);

    // Most important features should have higher scores
    let sorted_features: Vec<_> = importance_scores.iter().collect();
    // First feature should have highest importance in our mock implementation
    assert!(importance_scores.get(&0).unwrap() >= importance_scores.get(&7).unwrap());
}

#[tokio::test]
async fn test_attention_weight_analysis() {
    let config = TFTConfig {
        num_input_features: 6,
        hidden_dim: 24,
        num_attention_heads: 3,
        attention_dim: 8,
        num_quantiles: 5,
        quantiles: vec![0.1, 0.3, 0.5, 0.7, 0.9],
        dropout: 0.05,
        max_sequence_length: 10,
        prediction_horizons: vec![1, 2, 5],
        use_static_features: true,
        use_temporal_features: true,
        variable_selection_threshold: 0.1,
        attention_temperature: 1.2,
    };

    let mut model = MockTFTModel::new(config);
    let time_series = create_mock_time_series(8, 6);

    // Make prediction to compute attention weights
    let _ = model.forward(&time_series).await.unwrap();

    let attention_weights = model.get_attention_weights();
    assert!(!attention_weights.is_empty());

    // Attention weights should sum to 1
    let sum: f32 = attention_weights.iter().sum();
    assert!((sum - 1.0).abs() < 1e-5);

    // Recent positions should generally have higher attention
    // (This is implementation-specific and might not always hold)
    let num_weights = attention_weights.len();
    if num_weights > 2 {
        // Check that last few weights are reasonable
        assert!(attention_weights[num_weights - 1] >= 0.0);
        assert!(attention_weights[0] >= 0.0);
    }
}

#[tokio::test]
async fn test_prediction_interval_validity() {
    let config = TFTConfig {
        num_input_features: 4,
        hidden_dim: 20,
        num_attention_heads: 2,
        attention_dim: 10,
        num_quantiles: 9,
        quantiles: vec![0.05, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 0.95],
        dropout: 0.1,
        max_sequence_length: 12,
        prediction_horizons: vec![1, 3, 5],
        use_static_features: true,
        use_temporal_features: true,
        variable_selection_threshold: 0.05,
        attention_temperature: 1.0,
    };

    let mut model = MockTFTModel::new(config);
    let time_series = create_mock_time_series(10, 4);

    let prediction = model.forward(&time_series).await.unwrap();

    // Check prediction intervals are monotonically ordered
    let intervals = &prediction.prediction_intervals;

    if let Some(&(lower_50, upper_50)) = intervals.get("50%") {
        if let Some(&(lower_80, upper_80)) = intervals.get("80%") {
            // 80% interval should contain 50% interval
            assert!(lower_80 <= lower_50);
            assert!(upper_80 >= upper_50);
        }

        if let Some(&(lower_90, upper_90)) = intervals.get("90%") {
            // 90% interval should contain 50% interval
            assert!(lower_90 <= lower_50);
            assert!(upper_90 >= upper_50);
        }
    }

    // All intervals should have lower <= upper
    for (_, &(lower, upper)) in intervals {
        assert!(lower <= upper, "Interval bounds: {} <= {}", lower, upper);
    }
}

// Property-based tests using proptest
proptest! {
    #[test]
    fn test_tft_config_properties(
        num_input_features in 3..50_usize,
        hidden_dim in 16..256_usize,
        num_attention_heads in 1..8_usize,
        num_quantiles in 3..11_usize,
        max_sequence_length in 10..200_usize,
    ) {
        prop_assume!(hidden_dim % num_attention_heads == 0); // Hidden dim divisible by heads

        let quantiles: Vec<f64> = (1..=num_quantiles)
            .map(|i| (i as f64) / (num_quantiles + 1) as f64)
            .collect();

        let config = TFTConfig {
            num_input_features,
            hidden_dim,
            num_attention_heads,
            attention_dim: hidden_dim / num_attention_heads,
            num_quantiles,
            quantiles,
            dropout: 0.1,
            max_sequence_length,
            prediction_horizons: vec![1, 5, 10],
            use_static_features: true,
            use_temporal_features: true,
            variable_selection_threshold: 0.1,
            attention_temperature: 1.0,
        };

        let model = MockTFTModel::new(config.clone());
        prop_assert_eq!(model.config.num_input_features, num_input_features);
        prop_assert_eq!(model.config.hidden_dim, hidden_dim);
        prop_assert_eq!(model.config.num_attention_heads, num_attention_heads);
        prop_assert_eq!(model.config.num_quantiles, num_quantiles);
        prop_assert_eq!(model.config.quantiles.len(), num_quantiles);
    }

    #[test]
    fn test_quantile_ordering(
        quantiles in prop::collection::vec(0.01..0.99_f64, 3..10),
    ) {
        let rt = tokio::runtime::Runtime::new().unwrap();
        rt.block_on(async {
            let mut sorted_quantiles = quantiles.clone();
            sorted_quantiles.sort_by(|a, b| a.partial_cmp(b).unwrap());

            let config = TFTConfig {
                num_input_features: 5,
                hidden_dim: 32,
                num_attention_heads: 2,
                attention_dim: 16,
                num_quantiles: sorted_quantiles.len(),
                quantiles: sorted_quantiles.clone(),
                dropout: 0.1,
                max_sequence_length: 20,
                prediction_horizons: vec![1, 5],
                use_static_features: true,
                use_temporal_features: true,
                variable_selection_threshold: 0.1,
                attention_temperature: 1.0,
            };

            let mut model = MockTFTModel::new(config);
            let time_series = create_mock_time_series(10, 5);

            let prediction = model.forward(&time_series).await.unwrap();

            // Check that quantile predictions are monotonically increasing
            let mut quantile_values: Vec<(u8, f32)> = prediction.predictions.iter()
                .map(|(&k, &v)| (k, v))
                .collect();
            quantile_values.sort_by_key(|&(k, _)| k);

            for window in quantile_values.windows(2) {
                if let [(q1, v1), (q2, v2)] = window {
                    // Higher quantiles should generally have higher or equal values
                    // (This is a property of proper quantile predictions)
                    prop_assert!(q2 > q1); // Quantile keys are ordered
                    // Note: Values don't have to be strictly ordered due to our mock implementation
                    // but they should be finite
                    prop_assert!(v1.is_finite());
                    prop_assert!(v2.is_finite());
                }
            }
        });
    }
}