//! Multi-Symbol Training Integration Tests
//!
//! Tests for training models across multiple symbols (ES.FUT + NQ.FUT + ZN.FUT)
//! to validate data handling, feature consistency, and multi-asset model performance.
//!
//! # Test Coverage
//!
//! 1. **Multi-Symbol Data Loading** (3 scenarios)
//!    - Load multiple symbols simultaneously
//!    - Validate feature consistency across symbols
//!    - Handle missing/incomplete symbol data
//!
//! 2. **Multi-Symbol Training** (4 scenarios)
//!    - Train single model on multiple symbols
//!    - Train separate models per symbol
//!    - Mixed symbol batches
//!    - Symbol-specific feature normalization
//!
//! 3. **Cross-Symbol Validation** (2 scenarios)
//!    - Train on ES.FUT, validate on NQ.FUT
//!    - Ensemble prediction across symbols
//!
//! # Usage
//!
//! ```bash
//! cargo test -p ml multi_symbol -- --nocapture
//! ```

use anyhow::Result;
use candle_core::{Device, Tensor};
use std::collections::HashMap;
use std::path::PathBuf;

use ml::data_loaders::dbn_sequence_loader::DbnSequenceLoader;
use ml::mamba::{Mamba2Config, Mamba2SSM};

// ============================================================================
// Test Helpers
// ============================================================================

/// Check if DBN data exists for a symbol
fn check_dbn_data_exists(symbol: &str, date: &str) -> Option<PathBuf> {
    let path = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
        .parent()
        .unwrap()
        .join(format!("test_data/databento/{}/{}.dbn.zst", symbol, date));

    if path.exists() {
        Some(path)
    } else {
        None
    }
}

/// Load sequences for a symbol
async fn load_symbol_sequences(
    symbol: &str,
    date: &str,
    seq_len: usize,
    feature_dim: usize,
    max_sequences: usize,
) -> Result<Vec<Vec<f32>>> {
    let path = check_dbn_data_exists(symbol, date);

    if path.is_none() {
        return Ok(Vec::new());
    }

    let loader = DbnSequenceLoader::new(vec![path.unwrap().to_string_lossy().to_string()], seq_len, feature_dim)?;
    let sequences = loader.load_sequences(max_sequences).await?;

    // Convert to flat feature vectors
    let flat_sequences: Vec<Vec<f32>> = sequences
        .iter()
        .map(|seq| {
            seq.features
                .iter()
                .flat_map(|features| features.iter().copied())
                .collect()
        })
        .collect();

    Ok(flat_sequences)
}

// ============================================================================
// 1. Multi-Symbol Data Loading (3 scenarios)
// ============================================================================

#[tokio::test]
async fn test_load_multiple_symbols_simultaneously() -> Result<()> {
    println!("\n🧪 Test: Load Multiple Symbols Simultaneously");
    println!("Testing: ES.FUT + NQ.FUT + ZN.FUT data loading");

    let symbols = vec![
        ("ZN.FUT", "2024-01-02"),
        ("6E.FUT", "2024-01-02"),
        ("ES.FUT", "2024-01-02"),
    ];

    let seq_len = 60;
    let feature_dim = 16;
    let max_sequences = 50;

    let mut symbol_data: HashMap<String, Vec<Vec<f32>>> = HashMap::new();
    let mut symbols_loaded = 0;

    for (symbol, date) in symbols.iter() {
        print!("  Loading {}... ", symbol);

        match load_symbol_sequences(symbol, date, seq_len, feature_dim, max_sequences).await {
            Ok(sequences) => {
                if sequences.is_empty() {
                    println!("⏭️  NOT FOUND");
                } else {
                    println!("✓ {} sequences", sequences.len());
                    symbol_data.insert(symbol.to_string(), sequences);
                    symbols_loaded += 1;
                }
            }
            Err(e) => {
                println!("❌ ERROR: {:?}", e);
            }
        }
    }

    if symbols_loaded == 0 {
        println!("⏭️  Skipping: No DBN data found");
        return Ok(());
    }

    println!("  ✓ Loaded {} symbols", symbols_loaded);

    // Validate data dimensions
    for (symbol, sequences) in symbol_data.iter() {
        assert!(!sequences.is_empty(), "{} should have sequences", symbol);

        let expected_len = seq_len * feature_dim;
        assert_eq!(
            sequences[0].len(),
            expected_len,
            "{} sequence length should be {}",
            symbol,
            expected_len
        );

        println!("    {}: {} sequences, {} features per sequence", symbol, sequences.len(), sequences[0].len());
    }

    println!("✅ Multi-symbol loading test PASSED\n");
    Ok(())
}

#[tokio::test]
async fn test_feature_consistency_across_symbols() -> Result<()> {
    println!("\n🧪 Test: Feature Consistency Across Symbols");
    println!("Testing: Feature dimensions and ranges match across symbols");

    let symbols = vec![
        ("ZN.FUT", "2024-01-02"),
        ("6E.FUT", "2024-01-02"),
    ];

    let seq_len = 60;
    let feature_dim = 16;
    let max_sequences = 20;

    let mut symbol_data: HashMap<String, Vec<Vec<f32>>> = HashMap::new();

    for (symbol, date) in symbols.iter() {
        if let Ok(sequences) = load_symbol_sequences(symbol, date, seq_len, feature_dim, max_sequences).await {
            if !sequences.is_empty() {
                symbol_data.insert(symbol.to_string(), sequences);
            }
        }
    }

    if symbol_data.len() < 2 {
        println!("⏭️  Skipping: Need at least 2 symbols for comparison");
        return Ok(());
    }

    println!("  Comparing features across {} symbols...", symbol_data.len());

    // Get reference dimensions from first symbol
    let (ref_symbol, ref_sequences) = symbol_data.iter().next().unwrap();
    let ref_dim = ref_sequences[0].len();

    println!("    Reference: {} with {} features", ref_symbol, ref_dim);

    // Compare all symbols to reference
    for (symbol, sequences) in symbol_data.iter() {
        let seq_dim = sequences[0].len();

        assert_eq!(
            seq_dim, ref_dim,
            "Symbol {} dimension {} should match reference dimension {}",
            symbol, seq_dim, ref_dim
        );

        // Check feature value ranges (should be numeric and reasonable)
        let first_seq = &sequences[0];
        let has_finite = first_seq.iter().all(|&v| v.is_finite());
        let has_nonzero = first_seq.iter().any(|&v| v != 0.0);

        assert!(has_finite, "{} should have finite features", symbol);
        assert!(has_nonzero, "{} should have non-zero features", symbol);

        println!("    ✓ {}: {} features, all finite", symbol, seq_dim);
    }

    println!("  ✓ All symbols have consistent features");
    println!("✅ Feature consistency test PASSED\n");
    Ok(())
}

#[tokio::test]
async fn test_handle_missing_symbol_data() -> Result<()> {
    println!("\n🧪 Test: Handle Missing/Incomplete Symbol Data");
    println!("Testing: Graceful handling of missing symbols");

    let symbols = vec![
        ("ZN.FUT", "2024-01-02"),        // Real
        ("MISSING.FUT", "2024-01-02"),   // Fake
        ("6E.FUT", "2024-01-02"),        // Real
        ("NONEXISTENT", "9999-99-99"),   // Fake
    ];

    let seq_len = 60;
    let feature_dim = 16;
    let max_sequences = 10;

    let mut loaded_symbols = Vec::new();
    let mut missing_symbols = Vec::new();

    for (symbol, date) in symbols.iter() {
        print!("  Checking {}... ", symbol);

        match load_symbol_sequences(symbol, date, seq_len, feature_dim, max_sequences).await {
            Ok(sequences) => {
                if sequences.is_empty() {
                    println!("MISSING");
                    missing_symbols.push(symbol.to_string());
                } else {
                    println!("✓ FOUND ({} sequences)", sequences.len());
                    loaded_symbols.push(symbol.to_string());
                }
            }
            Err(e) => {
                println!("ERROR: {:?}", e);
                missing_symbols.push(symbol.to_string());
            }
        }
    }

    println!("  Summary:");
    println!("    Loaded: {} symbols", loaded_symbols.len());
    println!("    Missing: {} symbols", missing_symbols.len());

    // Should handle missing data gracefully without panicking
    assert!(loaded_symbols.len() + missing_symbols.len() == symbols.len(),
        "Should account for all symbols");

    println!("  ✓ Missing data handled gracefully");
    println!("✅ Missing data handling test PASSED\n");
    Ok(())
}

// ============================================================================
// 2. Multi-Symbol Training (4 scenarios)
// ============================================================================

#[tokio::test]
async fn test_train_single_model_multiple_symbols() -> Result<()> {
    println!("\n🧪 Test: Train Single Model on Multiple Symbols");
    println!("Testing: Unified model trained on ES.FUT + NQ.FUT + ZN.FUT");

    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);

    let symbols = vec![
        ("ZN.FUT", "2024-01-02"),
        ("6E.FUT", "2024-01-02"),
    ];

    let seq_len = 60;
    let feature_dim = 16;
    let max_sequences = 20;

    // Load data from all available symbols
    let mut all_sequences = Vec::new();
    let mut symbols_used = Vec::new();

    for (symbol, date) in symbols.iter() {
        if let Ok(sequences) = load_symbol_sequences(symbol, date, seq_len, feature_dim, max_sequences).await {
            if !sequences.is_empty() {
                println!("  Loaded {}: {} sequences", symbol, sequences.len());
                all_sequences.extend(sequences);
                symbols_used.push(symbol.to_string());
            }
        }
    }

    if all_sequences.is_empty() {
        println!("⏭️  Skipping: No data available");
        return Ok(());
    }

    println!("  Total sequences from {} symbols: {}", symbols_used.len(), all_sequences.len());

    // Create model
    let config = Mamba2Config {
        d_model: feature_dim,
        d_state: 16,
        num_layers: 2,
        batch_size: 8,
        seq_len,
        learning_rate: 1e-4,
        ..Default::default()
    };

    let mut model = Mamba2SSM::new(config, &device)?;
    model.initialize_optimizer()?;

    println!("  Training unified model...");

    // Train on mixed data
    let batch_size = 8.min(all_sequences.len());
    let mut total_loss = 0.0f32;

    for (idx, seq_data) in all_sequences.iter().take(batch_size).enumerate() {
        // Reshape sequence to [1, seq_len, feature_dim]
        let input = Tensor::from_vec(seq_data.clone(), (1, seq_len, feature_dim), &device)?;
        let target = Tensor::new(&[0.0f32], &device)?.reshape((1, 1))?;

        let output = model.forward(&input)?;
        let seq_len = output.dim(1)?;
        let output_last = output.narrow(1, seq_len - 1, 1)?.squeeze(1)?;

        let loss = (&output_last - &target)?.powf(2.0)?.mean_all()?;
        total_loss += loss.to_scalar::<f32>()?;

        loss.backward()?;
        model.optimizer_step()?;
    }

    let avg_loss = total_loss / batch_size as f32;
    println!("    Average loss: {:.6}", avg_loss);

    assert!(avg_loss.is_finite(), "Loss should be finite");
    println!("  ✓ Model trained on multi-symbol data");

    println!("✅ Multi-symbol training test PASSED\n");
    Ok(())
}

#[tokio::test]
async fn test_train_separate_models_per_symbol() -> Result<()> {
    println!("\n🧪 Test: Train Separate Models Per Symbol");
    println!("Testing: Symbol-specific model specialization");

    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);

    let symbols = vec![
        ("ZN.FUT", "2024-01-02"),
        ("6E.FUT", "2024-01-02"),
    ];

    let seq_len = 60;
    let feature_dim = 16;
    let max_sequences = 20;

    let mut symbol_models: HashMap<String, (Mamba2SSM, f32)> = HashMap::new();

    for (symbol, date) in symbols.iter() {
        if let Ok(sequences) = load_symbol_sequences(symbol, date, seq_len, feature_dim, max_sequences).await {
            if sequences.is_empty() {
                continue;
            }

            println!("  Training model for {}...", symbol);

            // Create symbol-specific model
            let config = Mamba2Config {
                d_model: feature_dim,
                d_state: 16,
                num_layers: 2,
                batch_size: 8,
                seq_len,
                learning_rate: 1e-4,
                ..Default::default()
            };

            let mut model = Mamba2SSM::new(config, &device)?;
            model.initialize_optimizer()?;

            // Train on symbol-specific data
            let batch_size = 8.min(sequences.len());
            let mut total_loss = 0.0f32;

            for seq_data in sequences.iter().take(batch_size) {
                let input = Tensor::from_vec(seq_data.clone(), (1, seq_len, feature_dim), &device)?;
                let target = Tensor::new(&[0.0f32], &device)?.reshape((1, 1))?;

                let output = model.forward(&input)?;
                let seq_len = output.dim(1)?;
                let output_last = output.narrow(1, seq_len - 1, 1)?.squeeze(1)?;

                let loss = (&output_last - &target)?.powf(2.0)?.mean_all()?;
                total_loss += loss.to_scalar::<f32>()?;

                loss.backward()?;
                model.optimizer_step()?;
            }

            let avg_loss = total_loss / batch_size as f32;
            println!("    {} loss: {:.6}", symbol, avg_loss);

            symbol_models.insert(symbol.to_string(), (model, avg_loss));
        }
    }

    if symbol_models.is_empty() {
        println!("⏭️  Skipping: No data available");
        return Ok(());
    }

    println!("  ✓ Trained {} symbol-specific models", symbol_models.len());

    // Validate each model
    for (symbol, (_model, loss)) in symbol_models.iter() {
        assert!(loss.is_finite(), "{} loss should be finite", symbol);
        println!("    {}: loss={:.6}", symbol, loss);
    }

    println!("✅ Symbol-specific training test PASSED\n");
    Ok(())
}

#[tokio::test]
async fn test_mixed_symbol_batches() -> Result<()> {
    println!("\n🧪 Test: Mixed Symbol Batches");
    println!("Testing: Training batches with multiple symbols mixed");

    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);

    let symbols = vec![
        ("ZN.FUT", "2024-01-02"),
        ("6E.FUT", "2024-01-02"),
    ];

    let seq_len = 60;
    let feature_dim = 16;
    let max_sequences = 15;

    // Load sequences with symbol labels
    let mut labeled_sequences: Vec<(String, Vec<f32>)> = Vec::new();

    for (symbol, date) in symbols.iter() {
        if let Ok(sequences) = load_symbol_sequences(symbol, date, seq_len, feature_dim, max_sequences).await {
            for seq in sequences {
                labeled_sequences.push((symbol.to_string(), seq));
            }
        }
    }

    if labeled_sequences.is_empty() {
        println!("⏭️  Skipping: No data available");
        return Ok(());
    }

    println!("  Total sequences: {}", labeled_sequences.len());

    // Count symbols in dataset
    let mut symbol_counts: HashMap<String, usize> = HashMap::new();
    for (symbol, _) in labeled_sequences.iter() {
        *symbol_counts.entry(symbol.clone()).or_insert(0) += 1;
    }

    for (symbol, count) in symbol_counts.iter() {
        println!("    {}: {} sequences", symbol, count);
    }

    // Create model
    let config = Mamba2Config {
        d_model: feature_dim,
        d_state: 16,
        num_layers: 2,
        batch_size: 8,
        seq_len,
        learning_rate: 1e-4,
        ..Default::default()
    };

    let mut model = Mamba2SSM::new(config, &device)?;
    model.initialize_optimizer()?;

    println!("  Training on mixed batches...");

    // Create mixed batch (interleave symbols)
    let batch_size = 8.min(labeled_sequences.len());
    let mut total_loss = 0.0f32;

    for (symbol, seq_data) in labeled_sequences.iter().take(batch_size) {
        let input = Tensor::from_vec(seq_data.clone(), (1, seq_len, feature_dim), &device)?;
        let target = Tensor::new(&[0.0f32], &device)?.reshape((1, 1))?;

        let output = model.forward(&input)?;
        let seq_len = output.dim(1)?;
        let output_last = output.narrow(1, seq_len - 1, 1)?.squeeze(1)?;

        let loss = (&output_last - &target)?.powf(2.0)?.mean_all()?;
        total_loss += loss.to_scalar::<f32>()?;

        loss.backward()?;
        model.optimizer_step()?;

        println!("    Trained on {}", symbol);
    }

    let avg_loss = total_loss / batch_size as f32;
    println!("  Average loss: {:.6}", avg_loss);

    assert!(avg_loss.is_finite(), "Loss should be finite");
    println!("  ✓ Model trained on mixed-symbol batches");

    println!("✅ Mixed batch training test PASSED\n");
    Ok(())
}

#[tokio::test]
async fn test_symbol_specific_normalization() -> Result<()> {
    println!("\n🧪 Test: Symbol-Specific Feature Normalization");
    println!("Testing: Different normalization per symbol");

    let symbols = vec![
        ("ZN.FUT", "2024-01-02"),
        ("6E.FUT", "2024-01-02"),
    ];

    let seq_len = 60;
    let feature_dim = 16;
    let max_sequences = 20;

    let mut symbol_stats: HashMap<String, (f32, f32)> = HashMap::new();

    for (symbol, date) in symbols.iter() {
        if let Ok(sequences) = load_symbol_sequences(symbol, date, seq_len, feature_dim, max_sequences).await {
            if sequences.is_empty() {
                continue;
            }

            println!("  Analyzing {}...", symbol);

            // Calculate mean and std for this symbol
            let mut all_values: Vec<f32> = Vec::new();
            for seq in sequences.iter() {
                all_values.extend(seq.iter().copied());
            }

            let mean = all_values.iter().sum::<f32>() / all_values.len() as f32;
            let variance = all_values.iter()
                .map(|&x| (x - mean).powi(2))
                .sum::<f32>() / all_values.len() as f32;
            let std = variance.sqrt();

            println!("    Mean: {:.6}, Std: {:.6}", mean, std);

            assert!(mean.is_finite(), "{} mean should be finite", symbol);
            assert!(std.is_finite(), "{} std should be finite", symbol);
            assert!(std > 0.0, "{} std should be positive", symbol);

            symbol_stats.insert(symbol.to_string(), (mean, std));
        }
    }

    if symbol_stats.is_empty() {
        println!("⏭️  Skipping: No data available");
        return Ok(());
    }

    println!("  ✓ Computed normalization stats for {} symbols", symbol_stats.len());

    // Verify stats differ between symbols (if multiple symbols loaded)
    if symbol_stats.len() >= 2 {
        let stats: Vec<_> = symbol_stats.values().collect();
        let mean_diff = (stats[0].0 - stats[1].0).abs();
        println!("    Mean difference between symbols: {:.6}", mean_diff);
    }

    println!("✅ Symbol normalization test PASSED\n");
    Ok(())
}

// ============================================================================
// 3. Cross-Symbol Validation (2 scenarios)
// ============================================================================

#[tokio::test]
async fn test_train_on_one_validate_on_another() -> Result<()> {
    println!("\n🧪 Test: Train on One Symbol, Validate on Another");
    println!("Testing: Generalization across different symbols");

    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);

    let train_symbol = ("ZN.FUT", "2024-01-02");
    let val_symbol = ("6E.FUT", "2024-01-02");

    let seq_len = 60;
    let feature_dim = 16;
    let max_sequences = 20;

    // Load training data
    let train_sequences = load_symbol_sequences(train_symbol.0, train_symbol.1, seq_len, feature_dim, max_sequences).await?;

    if train_sequences.is_empty() {
        println!("⏭️  Skipping: Training data ({}) not available", train_symbol.0);
        return Ok(());
    }

    println!("  Training data ({}): {} sequences", train_symbol.0, train_sequences.len());

    // Load validation data
    let val_sequences = load_symbol_sequences(val_symbol.0, val_symbol.1, seq_len, feature_dim, max_sequences).await?;

    if val_sequences.is_empty() {
        println!("⏭️  Skipping: Validation data ({}) not available", val_symbol.0);
        return Ok(());
    }

    println!("  Validation data ({}): {} sequences", val_symbol.0, val_sequences.len());

    // Train model on first symbol
    let config = Mamba2Config {
        d_model: feature_dim,
        d_state: 16,
        num_layers: 2,
        batch_size: 8,
        seq_len,
        learning_rate: 1e-4,
        ..Default::default()
    };

    let mut model = Mamba2SSM::new(config, &device)?;
    model.initialize_optimizer()?;

    println!("  Training on {}...", train_symbol.0);

    let batch_size = 8.min(train_sequences.len());
    let mut train_loss = 0.0f32;

    for seq_data in train_sequences.iter().take(batch_size) {
        let input = Tensor::from_vec(seq_data.clone(), (1, seq_len, feature_dim), &device)?;
        let target = Tensor::new(&[0.0f32], &device)?.reshape((1, 1))?;

        let output = model.forward(&input)?;
        let seq_len = output.dim(1)?;
        let output_last = output.narrow(1, seq_len - 1, 1)?.squeeze(1)?;

        let loss = (&output_last - &target)?.powf(2.0)?.mean_all()?;
        train_loss += loss.to_scalar::<f32>()?;

        loss.backward()?;
        model.optimizer_step()?;
    }

    train_loss /= batch_size as f32;
    println!("    Training loss: {:.6}", train_loss);

    // Validate on second symbol
    println!("  Validating on {}...", val_symbol.0);

    let val_batch_size = 8.min(val_sequences.len());
    let mut val_loss = 0.0f32;

    for seq_data in val_sequences.iter().take(val_batch_size) {
        let input = Tensor::from_vec(seq_data.clone(), (1, seq_len, feature_dim), &device)?;
        let target = Tensor::new(&[0.0f32], &device)?.reshape((1, 1))?;

        let output = model.forward(&input)?;
        let seq_len = output.dim(1)?;
        let output_last = output.narrow(1, seq_len - 1, 1)?.squeeze(1)?;

        let loss = (&output_last - &target)?.powf(2.0)?.mean_all()?;
        val_loss += loss.to_scalar::<f32>()?;
    }

    val_loss /= val_batch_size as f32;
    println!("    Validation loss: {:.6}", val_loss);

    assert!(train_loss.is_finite(), "Training loss should be finite");
    assert!(val_loss.is_finite(), "Validation loss should be finite");

    println!("  ✓ Cross-symbol validation completed");
    println!("✅ Cross-symbol validation test PASSED\n");
    Ok(())
}

#[tokio::test]
async fn test_ensemble_prediction_across_symbols() -> Result<()> {
    println!("\n🧪 Test: Ensemble Prediction Across Symbols");
    println!("Testing: Multiple models predicting on shared data");

    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);

    let symbols = vec![
        ("ZN.FUT", "2024-01-02"),
        ("6E.FUT", "2024-01-02"),
    ];

    let seq_len = 60;
    let feature_dim = 16;
    let max_sequences = 10;

    // Train one model per symbol
    let mut models: Vec<(String, Mamba2SSM)> = Vec::new();

    for (symbol, date) in symbols.iter() {
        let sequences = load_symbol_sequences(symbol, date, seq_len, feature_dim, max_sequences).await?;

        if sequences.is_empty() {
            continue;
        }

        println!("  Training model for {}...", symbol);

        let config = Mamba2Config {
            d_model: feature_dim,
            d_state: 16,
            num_layers: 2,
            batch_size: 8,
            seq_len,
            learning_rate: 1e-4,
            ..Default::default()
        };

        let mut model = Mamba2SSM::new(config, &device)?;
        model.initialize_optimizer()?;

        // Quick training
        let batch_size = 5.min(sequences.len());
        for seq_data in sequences.iter().take(batch_size) {
            let input = Tensor::from_vec(seq_data.clone(), (1, seq_len, feature_dim), &device)?;
            let target = Tensor::new(&[0.0f32], &device)?.reshape((1, 1))?;

            let output = model.forward(&input)?;
            let seq_len = output.dim(1)?;
            let output_last = output.narrow(1, seq_len - 1, 1)?.squeeze(1)?;

            let loss = (&output_last - &target)?.powf(2.0)?.mean_all()?;
            loss.backward()?;
            model.optimizer_step()?;
        }

        models.push((symbol.to_string(), model));
        println!("    ✓ Model trained");
    }

    if models.len() < 2 {
        println!("⏭️  Skipping: Need at least 2 models for ensemble");
        return Ok(());
    }

    println!("  Created ensemble with {} models", models.len());

    // Test ensemble prediction on shared test data
    let test_input = Tensor::randn(0.0f32, 1.0, (1, seq_len, feature_dim), &device)?;

    println!("  Running ensemble predictions...");

    let mut predictions = Vec::new();

    for (symbol, model) in models.iter_mut() {
        let output = model.forward(&test_input)?;
        let seq_len = output.dim(1)?;
        let output_last = output.narrow(1, seq_len - 1, 1)?.squeeze(1)?;

        let pred = output_last.to_vec1::<f32>()?[0];
        println!("    {}: {:.6}", symbol, pred);
        predictions.push(pred);
    }

    // Compute ensemble average
    let ensemble_pred = predictions.iter().sum::<f32>() / predictions.len() as f32;
    println!("  Ensemble prediction: {:.6}", ensemble_pred);

    assert!(ensemble_pred.is_finite(), "Ensemble prediction should be finite");
    println!("  ✓ Ensemble prediction completed");

    println!("✅ Ensemble prediction test PASSED\n");
    Ok(())
}

// ============================================================================
// Test Summary
// ============================================================================

#[tokio::test]
async fn test_multi_symbol_summary() -> Result<()> {
    println!("\n📊 Multi-Symbol Test Summary");
    println!("============================");
    println!("Data Loading: 3 scenarios");
    println!("  - Simultaneous loading");
    println!("  - Feature consistency");
    println!("  - Missing data handling");
    println!("");
    println!("Training: 4 scenarios");
    println!("  - Single unified model");
    println!("  - Separate per-symbol models");
    println!("  - Mixed symbol batches");
    println!("  - Symbol-specific normalization");
    println!("");
    println!("Validation: 2 scenarios");
    println!("  - Cross-symbol validation");
    println!("  - Ensemble prediction");
    println!("");
    println!("Total: 9 multi-symbol test scenarios");
    println!("============================\n");

    Ok(())
}