foxhunt/ml/tests/dqn_parquet_loading_test.rs

//! DQN Parquet Loading Test
//!
//! Validates that the DQN adapter can correctly load parquet files
//! and extract 225-feature vectors.

use ml::hyperopt::adapters::dqn::DQNTrainer;
use std::path::PathBuf;

// ============================================================================
// Component 4: Inference Engine
// ============================================================================

use std::time::Instant;

/// Result of a single DQN inference
#[derive(Debug, Clone)]
struct InferenceResult {
    action: usize,      // 0=BUY, 1=SELL, 2=HOLD
    q_values: [f32; 3], // Q-value for each action
    latency_us: u64,    // Microseconds for this inference
}

/// Run DQN inference on all feature vectors with progress tracking
///
/// # Arguments
/// * `network` - QNetwork for inference (the underlying network from DQNAgent)
/// * `features` - 225-dimensional feature vectors
///
/// # Returns
/// Vector of inference results (action, Q-values, latency per bar)
///
/// # Notes
/// - Uses single-sample batches (shape [1, 225]) for inference
/// - Handles NaN/Inf gracefully by logging warnings and skipping bars
/// - Tracks latency per inference in microseconds
/// - Progress bar shows real-time inference speed
fn run_inference(
    network: &ml::dqn::network::QNetwork,
    features: Vec<[f64; 225]>,
) -> Result<Vec<InferenceResult>, anyhow::Error> {
    let total_bars = features.len();
    if total_bars == 0 {
        anyhow::bail!("No feature vectors provided for inference");
    }

    println!("\n🔍 Starting DQN Inference");
    println!("  Total bars to process: {}", total_bars);

    let mut results = Vec::with_capacity(total_bars);
    let mut total_inference_time_us = 0u64;
    let mut skipped_bars = 0usize;
    let start_time = Instant::now();
    let mut last_progress_update = Instant::now();

    // Run inference for each feature vector
    for (i, feature_vec) in features.iter().enumerate() {
        // Start timer for this inference
        let timer = Instant::now();

        // Convert f64 features to f32 for DQN network
        let state_f32: Vec<f32> = feature_vec.iter().map(|&x| x as f32).collect();

        // Forward pass to get Q-values
        let q_values_result = network.forward(&state_f32);

        // Handle forward pass errors
        let q_values_vec = match q_values_result {
            Ok(qv) => qv,
            Err(e) => {
                if skipped_bars < 10 {
                    eprintln!("  ⚠️  Bar {}: Forward pass failed: {}. Skipping.", i, e);
                }
                skipped_bars += 1;
                continue;
            },
        };

        // Ensure we have exactly 3 Q-values (BUY, SELL, HOLD)
        if q_values_vec.len() != 3 {
            if skipped_bars < 10 {
                eprintln!(
                    "  ⚠️  Bar {}: Expected 3 Q-values, got {}. Skipping.",
                    i,
                    q_values_vec.len()
                );
            }
            skipped_bars += 1;
            continue;
        }

        let q_values: [f32; 3] = [q_values_vec[0], q_values_vec[1], q_values_vec[2]];

        // Check for NaN/Inf in Q-values
        if q_values.iter().any(|&q| !q.is_finite()) {
            if skipped_bars < 10 {
                eprintln!(
                    "  ⚠️  Bar {}: Q-values contain NaN/Inf, skipping. Q-values: {:?}",
                    i, q_values
                );
            }
            skipped_bars += 1;
            continue;
        }

        // Select action with argmax(q_values)
        let action = q_values
            .iter()
            .enumerate()
            .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
            .map(|(idx, _)| idx)
            .unwrap_or(2); // Default to HOLD if comparison fails

        // Calculate latency for this inference
        let latency_us = timer.elapsed().as_micros() as u64;
        total_inference_time_us += latency_us;

        // Store result
        results.push(InferenceResult {
            action,
            q_values,
            latency_us,
        });

        // Update progress every 1 second or every 10% completion
        let should_update = last_progress_update.elapsed().as_secs() >= 1
            || (i + 1) % (total_bars / 10).max(1) == 0
            || i == 0
            || i == total_bars - 1;

        if should_update {
            let elapsed_sec = start_time.elapsed().as_secs_f64();
            let avg_speed = if elapsed_sec > 0.0 {
                (i + 1) as f64 / elapsed_sec
            } else {
                0.0
            };
            let progress_pct = ((i + 1) as f64 / total_bars as f64) * 100.0;
            println!(
                "  Progress: {}/{} ({:.1}%) | Speed: {:.1} bars/sec | Skipped: {}",
                i + 1,
                total_bars,
                progress_pct,
                avg_speed,
                skipped_bars
            );
            last_progress_update = Instant::now();
        }
    }

    println!("\n✅ Inference complete");

    // Calculate summary statistics
    let processed_bars = total_bars - skipped_bars;
    let total_time_sec = start_time.elapsed().as_secs_f64();
    let avg_latency_us = if processed_bars > 0 {
        total_inference_time_us / processed_bars as u64
    } else {
        0
    };
    let avg_speed = if total_time_sec > 0.0 {
        processed_bars as f64 / total_time_sec
    } else {
        0.0
    };

    // Log summary with detailed metrics
    println!("\n========================================");
    println!("📊 Inference Summary:");
    println!("========================================");
    println!("  Total bars: {}", total_bars);
    println!("  Processed: {}", processed_bars);
    println!("  Skipped (NaN/Inf/errors): {}", skipped_bars);
    println!(
        "  Skip rate: {:.2}%",
        (skipped_bars as f64 / total_bars as f64) * 100.0
    );
    println!("  Total time: {:.2}s", total_time_sec);
    println!(
        "  Average latency: {}μs ({:.2}ms)",
        avg_latency_us,
        avg_latency_us as f64 / 1000.0
    );
    println!("  Average speed: {:.1} bars/sec", avg_speed);
    println!("========================================");

    // Validate results
    if results.is_empty() {
        anyhow::bail!(
            "All {} inference attempts failed (likely NaN/Inf in Q-values or network errors)",
            total_bars
        );
    }

    // Log action distribution
    let mut action_counts = [0usize; 3];
    for result in &results {
        action_counts[result.action] += 1;
    }
    println!("\n📈 Action Distribution:");
    println!(
        "  BUY:  {} ({:.1}%)",
        action_counts[0],
        (action_counts[0] as f64 / processed_bars as f64) * 100.0
    );
    println!(
        "  SELL: {} ({:.1}%)",
        action_counts[1],
        (action_counts[1] as f64 / processed_bars as f64) * 100.0
    );
    println!(
        "  HOLD: {} ({:.1}%)",
        action_counts[2],
        (action_counts[2] as f64 / processed_bars as f64) * 100.0
    );

    if skipped_bars > 10 {
        println!(
            "\n⚠️  Warning: {} additional errors were silenced (only first 10 shown)",
            skipped_bars - 10
        );
    }

    Ok(results)
}

// ============================================================================
// End Component 4
// ============================================================================

// ============================================================================
// Component 4 Test: Inference Engine Validation
// ============================================================================

#[test]
fn test_inference_engine_with_mock_network() {
    use ml::dqn::network::{QNetwork, QNetworkConfig};

    // Create a simple DQN network for testing
    let config = QNetworkConfig {
        state_dim: 225,
        num_actions: 3,
        hidden_dims: vec![64, 32],
        learning_rate: 0.001,
        epsilon_start: 0.0, // No exploration for deterministic testing
        epsilon_end: 0.0,
        epsilon_decay: 1.0,
        target_update_freq: 1000,
        dropout_prob: 0.0, // No dropout for testing
        use_gpu: false,    // CPU only for testing
    };

    let network = QNetwork::new(config).expect("Failed to create QNetwork");

    // Create mock feature vectors (10 bars with 225 features each)
    let features: Vec<[f64; 225]> = (0..10)
        .map(|i| {
            let mut feature_vec = [0.0f64; 225];
            // Fill with some deterministic values
            for (j, val) in feature_vec.iter_mut().enumerate() {
                *val = (i as f64 + j as f64 * 0.01).sin();
            }
            feature_vec
        })
        .collect();

    // Run inference
    let results = run_inference(&network, features);

    // Validate results
    assert!(results.is_ok(), "Inference should succeed");

    let inference_results = results.unwrap();
    assert_eq!(
        inference_results.len(),
        10,
        "Should have 10 inference results"
    );

    // Validate each result
    for (i, result) in inference_results.iter().enumerate() {
        assert!(
            result.action < 3,
            "Action {} should be 0 (BUY), 1 (SELL), or 2 (HOLD)",
            result.action
        );

        // Q-values should be finite
        for (j, &q) in result.q_values.iter().enumerate() {
            assert!(
                q.is_finite(),
                "Q-value[{}] at bar {} should be finite, got {}",
                j,
                i,
                q
            );
        }

        // Latency should be reasonable (< 10ms per inference on CPU)
        assert!(
            result.latency_us < 10_000,
            "Latency at bar {} should be < 10ms, got {}μs",
            i,
            result.latency_us
        );
    }

    println!("✅ Inference engine test passed!");
}

#[test]
fn test_inference_engine_handles_empty_features() {
    use ml::dqn::network::{QNetwork, QNetworkConfig};

    let config = QNetworkConfig {
        state_dim: 225,
        num_actions: 3,
        hidden_dims: vec![64, 32],
        learning_rate: 0.001,
        epsilon_start: 0.0,
        epsilon_end: 0.0,
        epsilon_decay: 1.0,
        target_update_freq: 1000,
        dropout_prob: 0.0,
        use_gpu: false,
    };

    let network = QNetwork::new(config).expect("Failed to create QNetwork");

    // Empty feature vector
    let features: Vec<[f64; 225]> = vec![];

    // Run inference
    let results = run_inference(&network, features);

    // Should fail with empty input
    assert!(
        results.is_err(),
        "Inference should fail with empty feature vector"
    );

    let error_msg = format!("{:?}", results.unwrap_err());
    assert!(
        error_msg.contains("No feature vectors"),
        "Error should mention empty input"
    );

    println!("✅ Empty features test passed!");
}

// ============================================================================
// End Component 4 Tests
// ============================================================================

#[test]
fn test_dqn_parquet_loading_small_file() {
    // Test with small parquet file
    let test_data_dir = PathBuf::from("test_data");

    // Verify test data exists
    let parquet_file = test_data_dir.join("ES_FUT_small.parquet");
    assert!(
        parquet_file.exists(),
        "Test parquet file not found: {:?}",
        parquet_file
    );

    // Create DQN trainer pointing to directory with parquet file
    let trainer_result = DQNTrainer::new(&test_data_dir, 5);

    assert!(
        trainer_result.is_ok(),
        "Failed to create DQN trainer: {:?}",
        trainer_result.err()
    );

    let trainer = trainer_result.unwrap();

    // Test that the trainer can detect parquet files
    // This would call load_training_data() internally
    // For now, we're just validating construction works

    println!("✓ DQN trainer created successfully with parquet data directory");
}

#[test]
fn test_dqn_parquet_file_detection() {
    // Test that DQN trainer prefers parquet over DBN when both exist
    let test_data_dir = PathBuf::from("test_data");

    assert!(
        test_data_dir.exists(),
        "Test data directory not found: {:?}",
        test_data_dir
    );

    // Count parquet files
    let parquet_count = std::fs::read_dir(&test_data_dir)
        .unwrap()
        .filter_map(|entry| entry.ok())
        .filter(|entry| entry.path().extension().and_then(|s| s.to_str()) == Some("parquet"))
        .count();

    assert!(
        parquet_count > 0,
        "No parquet files found in test_data directory"
    );

    println!(
        "✓ Found {} parquet file(s) in test data directory",
        parquet_count
    );
}

#[test]
fn test_dqn_requires_parquet_or_dbn() {
    // Test that DQN trainer fails gracefully when no data files exist
    use tempfile::TempDir;

    let temp_dir = TempDir::new().unwrap();
    let empty_dir = temp_dir.path();

    // Create DQN trainer with empty directory
    let trainer = DQNTrainer::new(empty_dir, 5);

    // Should succeed in creating trainer (validation happens at training time)
    assert!(
        trainer.is_ok(),
        "DQN trainer should accept empty directory at construction"
    );

    println!("✓ DQN trainer construction doesn't require immediate file validation");
}

#[test]
fn test_dqn_auto_detects_parquet() {
    // Test that DQN adapter auto-detects parquet files
    use ml::hyperopt::adapters::dqn::DQNParams;
    use ml::hyperopt::traits::HyperparameterOptimizable;

    let test_data_dir = PathBuf::from("test_data");

    // Create trainer pointing to directory with parquet files
    let mut trainer = DQNTrainer::new(&test_data_dir, 2).unwrap();

    // Create test parameters
    let params = DQNParams {
        learning_rate: 0.001,
        batch_size: 64,
        gamma: 0.99,
        epsilon_decay: 0.995,
        buffer_size: 10_000,
    };

    // This should use train_from_parquet() internally since parquet files exist
    // Note: This will actually try to train, so we expect it to work or fail with
    // training errors, not "No DBN files found"
    let result = trainer.train_with_params(params);

    // We expect either success OR a training error (not "No DBN files found")
    match result {
        Ok(metrics) => {
            println!("✓ DQN trained successfully with parquet data");
            println!("  Train loss: {:.6}", metrics.train_loss);
            assert!(metrics.train_loss.is_finite(), "Loss should be finite");
        },
        Err(e) => {
            let error_msg = format!("{:?}", e);
            // Should NOT contain "No DBN files found"
            assert!(
                !error_msg.contains("No DBN files found"),
                "DQN should use parquet files, not DBN. Error: {}",
                error_msg
            );
            println!("✓ DQN attempted parquet training (got training error, not DBN error)");
        },
    }
}