foxhunt/ml/tests/dqn_training_loop_integration_test.rs

//! DQN Training Loop Integration Tests
//!
//! Wave 10-A18: Tests to expose the dual reward system bug and validate the fix.
//!
//! **Bug Description**:
//! The main training loop `train_with_data_full_loop()` uses simple match-based rewards
//! (HOLD = -0.0001 fixed) instead of the sophisticated RewardFunction with portfolio
//! tracking, movement thresholds, and diversity penalties. This causes 100% HOLD bias.
//!
//! **These tests**:
//! 1. Expose the bug by showing HOLD is learned preferentially
//! 2. Validate that RewardFunction (when used) produces proper diversity
//! 3. Verify target network updates don't interfere with learning
//! 4. Test epsilon decay doesn't force premature exploitation

use anyhow::Result;
use ml::dqn::{Experience, TradingAction, TradingState};
use ml::trainers::dqn::{DQNHyperparameters, DQNTrainer};
use rust_decimal::Decimal;
use std::collections::HashMap;

/// Helper: Create minimal test hyperparameters
fn create_test_hyperparams() -> DQNHyperparameters {
    DQNHyperparameters {
        learning_rate: 0.001,
        batch_size: 4,
        gamma: 0.99,
        epsilon_start: 0.1, // Low epsilon for deterministic testing
        epsilon_end: 0.01,
        epsilon_decay: 0.99,
        buffer_size: 1000,
        min_replay_size: 10,
        epochs: 1,
        checkpoint_frequency: 100,
        early_stopping_enabled: false,
        q_value_floor: 0.5,
        min_loss_improvement_pct: 2.0,
        plateau_window: 5,
        min_epochs_before_stopping: 50,
        hold_penalty: -0.001,
        use_huber_loss: true,
        huber_delta: 1.0,
        use_double_dqn: true,
        gradient_clip_norm: Some(10.0),
        hold_penalty_weight: 0.01,
        movement_threshold: 0.02,
    }
}

/// Helper: Create synthetic training data with clear patterns
///
/// Pattern: Price increases by 5 points per step (5900 → 5905 → 5910 → ...)
/// Optimal policy: BUY when price going up, SELL when going down, HOLD when flat
fn create_synthetic_uptrend_data() -> Vec<([f64; 225], Vec<f64>)> {
    let mut data = Vec::new();
    let base_price = 5900.0;

    for i in 0..100 {
        let current_price = base_price + (i as f64 * 5.0);
        let next_price = current_price + 5.0;

        // Create 225-dim feature vector (Wave C + Wave D)
        let mut features = [0.0; 225];
        features[0] = current_price; // open
        features[1] = current_price + 2.0; // high
        features[2] = current_price - 1.0; // low
        features[3] = current_price; // close (most important for reward)

        // Fill remaining features with small random values
        for j in 4..225 {
            features[j] = (i as f64 * 0.01) + (j as f64 * 0.001);
        }

        // Target: [current_close, next_close]
        let target = vec![current_price, next_price];

        data.push((features, target));
    }

    data
}

/// Helper: Create synthetic downtrend data
fn create_synthetic_downtrend_data() -> Vec<([f64; 225], Vec<f64>)> {
    let mut data = Vec::new();
    let base_price = 6000.0;

    for i in 0..100 {
        let current_price = base_price - (i as f64 * 5.0);
        let next_price = current_price - 5.0;

        let mut features = [0.0; 225];
        features[0] = current_price;
        features[1] = current_price + 1.0;
        features[2] = current_price - 2.0;
        features[3] = current_price;

        for j in 4..225 {
            features[j] = (i as f64 * 0.01) + (j as f64 * 0.001);
        }

        let target = vec![current_price, next_price];
        data.push((features, target));
    }

    data
}

/// Helper: Create synthetic flat market data
fn create_synthetic_flat_data() -> Vec<([f64; 225], Vec<f64>)> {
    let mut data = Vec::new();
    let base_price = 5950.0;

    for i in 0..100 {
        let current_price = base_price; // No price movement
        let next_price = base_price;

        let mut features = [0.0; 225];
        features[0] = current_price;
        features[1] = current_price;
        features[2] = current_price;
        features[3] = current_price;

        for j in 4..225 {
            features[j] = (i as f64 * 0.01) + (j as f64 * 0.001);
        }

        let target = vec![current_price, next_price];
        data.push((features, target));
    }

    data
}

#[tokio::test]
async fn test_full_training_loop_learns_uptrend_policy() -> Result<()> {
    // **TEST OBJECTIVE**: Verify that after training on uptrend data, the agent
    // learns to prefer BUY actions over HOLD.
    //
    // **EXPECTED (with correct RewardFunction)**:
    // - BUY actions should be > 30% (agent learns to buy in uptrends)
    // - HOLD actions should be < 70% (agent avoids holding when profitable to buy)
    //
    // **CURRENT BUG (with simple match rewards)**:
    // - HOLD actions ~100% (agent learns HOLD is safest due to tiny -0.0001 penalty)

    let hyperparams = create_test_hyperparams();
    let mut trainer = DQNTrainer::new(hyperparams)?;

    // Generate 100 samples of uptrend data
    let training_data = create_synthetic_uptrend_data();

    // Train for 10 steps
    let mut action_counts = HashMap::new();
    action_counts.insert(TradingAction::Buy, 0);
    action_counts.insert(TradingAction::Sell, 0);
    action_counts.insert(TradingAction::Hold, 0);

    // Simulate 10 training steps
    for (features, _target) in training_data.iter().take(10) {
        // Convert to trading state
        let close_price = Decimal::try_from(features[3]).unwrap_or(Decimal::ZERO);
        let state = trainer.feature_vector_to_state(features, Some(close_price))?;

        // Select action
        let action = trainer.select_action(&state).await?;
        *action_counts.entry(action).or_insert(0) += 1;
    }

    let total_actions: usize = action_counts.values().sum();
    let buy_pct = (*action_counts.get(&TradingAction::Buy).unwrap_or(&0) as f64
        / total_actions as f64)
        * 100.0;
    let hold_pct = (*action_counts.get(&TradingAction::Hold).unwrap_or(&0) as f64
        / total_actions as f64)
        * 100.0;

    println!("Uptrend Policy Test:");
    println!("  BUY: {:.1}%", buy_pct);
    println!(
        "  SELL: {:.1}%",
        (*action_counts.get(&TradingAction::Sell).unwrap_or(&0) as f64 / total_actions as f64)
            * 100.0
    );
    println!("  HOLD: {:.1}%", hold_pct);

    // **ASSERTION REVEALS BUG**:
    // With simple rewards: This test will FAIL (HOLD ~100%)
    // With RewardFunction: This test will PASS (BUY > 30%, HOLD < 70%)
    assert!(
        hold_pct < 90.0,
        "HOLD bias detected: {:.1}% HOLD actions (expected < 90%). Bug: Simple match rewards favor HOLD.",
        hold_pct
    );

    Ok(())
}

#[tokio::test]
async fn test_target_network_stabilizes_learning() -> Result<()> {
    // **TEST OBJECTIVE**: Verify that target network updates don't interfere with learning.
    //
    // **METHOD**: Train for 20 steps and track Q-value stability. Target network should
    // reduce oscillations compared to no target network.

    let mut hyperparams = create_test_hyperparams();
    hyperparams.min_replay_size = 5; // Allow training after 5 experiences

    let mut trainer = DQNTrainer::new(hyperparams)?;

    let training_data = create_synthetic_uptrend_data();
    let mut q_value_history = Vec::new();

    // Populate replay buffer with 10 experiences
    for (features, target) in training_data.iter().take(10) {
        let close_price = Decimal::try_from(features[3]).unwrap_or(Decimal::ZERO);
        let state = trainer.feature_vector_to_state(features, Some(close_price))?;
        let action = TradingAction::Buy; // Fixed action for consistency

        let next_close = if target.len() >= 2 {
            target[1]
        } else {
            features[3]
        };
        let next_close_price = Decimal::try_from(next_close).unwrap_or(Decimal::ZERO);
        let next_state = trainer.feature_vector_to_state(features, Some(next_close_price))?;

        let experience = Experience::new(
            state.to_vector(),
            action.to_int(),
            0.5, // Fixed reward
            next_state.to_vector(),
            false,
        );

        trainer.store_experience(experience).await?;
    }

    // Perform 10 training steps and track Q-values
    for _ in 0..10 {
        if trainer.can_train().await? {
            let (_loss, q_value, _grad_norm) = trainer.train_step().await?;
            q_value_history.push(q_value);
        }
    }

    println!("Target Network Stability Test:");
    println!("  Q-value history: {:?}", q_value_history);

    // Calculate Q-value variance (should be low if target network stabilizes)
    if q_value_history.len() > 1 {
        let mean = q_value_history.iter().sum::<f64>() / q_value_history.len() as f64;
        let variance = q_value_history
            .iter()
            .map(|q| (q - mean).powi(2))
            .sum::<f64>()
            / q_value_history.len() as f64;
        let std = variance.sqrt();

        println!("  Q-value std: {:.4}", std);

        // Target network should keep std reasonable (< 10.0)
        assert!(
            std < 10.0,
            "Q-value oscillation too high: std={:.4} (expected < 10.0). Target network may not be stabilizing.",
            std
        );
    }

    Ok(())
}

#[tokio::test]
async fn test_epsilon_decay_allows_exploration() -> Result<()> {
    // **TEST OBJECTIVE**: Verify that epsilon decay rate allows sufficient exploration.
    //
    // **METHOD**: Track epsilon values over 100 steps. With decay=0.995, epsilon should
    // decay slowly enough to explore for at least 50 steps.

    let mut hyperparams = create_test_hyperparams();
    hyperparams.epsilon_start = 1.0;
    hyperparams.epsilon_decay = 0.995;
    hyperparams.epsilon_end = 0.01;

    let trainer = DQNTrainer::new(hyperparams)?;

    // Simulate epsilon decay over 100 steps
    let initial_epsilon = trainer.get_epsilon().await?;
    println!("Epsilon Decay Test:");
    println!("  Initial epsilon: {:.4}", initial_epsilon);

    // Check epsilon after 50 steps (simulate by calculating)
    let epsilon_after_50 = initial_epsilon * 0.995_f32.powi(50);
    println!("  Epsilon after 50 steps: {:.4}", epsilon_after_50);

    // Epsilon should still be > 0.5 after 50 steps for good exploration
    assert!(
        epsilon_after_50 > 0.5,
        "Epsilon decays too fast: {:.4} after 50 steps (expected > 0.5). Increase epsilon_decay closer to 1.0.",
        epsilon_after_50
    );

    Ok(())
}

#[tokio::test]
async fn test_reward_function_diversity_penalty() -> Result<()> {
    // **TEST OBJECTIVE**: Verify that RewardFunction applies diversity penalty correctly.
    //
    // **METHOD**: Create a scenario where agent repeatedly selects HOLD. RewardFunction
    // should apply increasing diversity penalties.
    //
    // **NOTE**: This test directly uses RewardFunction, NOT the training loop, to verify
    // the correct implementation exists (even if unused in production).

    use ml::dqn::reward::{RewardConfig, RewardFunction};

    let reward_config = RewardConfig {
        pnl_weight: Decimal::ONE,
        risk_weight: Decimal::try_from(0.1).unwrap_or(Decimal::ZERO),
        cost_weight: Decimal::try_from(0.05).unwrap_or(Decimal::ZERO),
        hold_reward: Decimal::try_from(0.001).unwrap_or(Decimal::ZERO),
        movement_threshold: Decimal::try_from(0.02).unwrap_or(Decimal::ZERO),
        hold_penalty_weight: Decimal::try_from(0.01).unwrap_or(Decimal::ZERO),
        diversity_weight: Decimal::try_from(-0.1).unwrap_or(Decimal::ZERO),
    };

    let reward_fn = RewardFunction::new(reward_config);

    // Create state with flat prices (no movement)
    let state = TradingState {
        open: 5900.0,
        high: 5900.0,
        low: 5900.0,
        close: 5900.0,
        volume: 1000.0,
        technical_indicators: vec![0.0; 16],
        microstructure_features: vec![0.0; 16],
        portfolio_features: vec![0.0; 16],
        tick_imbalance: 0.0,
        order_flow_imbalance: 0.0,
        bid_ask_spread: 0.01,
    };

    let next_state = TradingState {
        close: 5900.0, // No price change
        ..state.clone()
    };

    // Test: Repeated HOLD actions should accumulate diversity penalty
    let recent_actions_uniform = vec![TradingAction::Buy, TradingAction::Sell, TradingAction::Hold];
    let recent_actions_biased = vec![TradingAction::Hold; 10]; // 10x HOLD

    let reward_uniform = reward_fn.calculate_reward(
        TradingAction::Hold,
        &state,
        &next_state,
        &recent_actions_uniform,
    )?;

    let reward_biased = reward_fn.calculate_reward(
        TradingAction::Hold,
        &state,
        &next_state,
        &recent_actions_biased,
    )?;

    println!("Diversity Penalty Test:");
    println!("  Reward (uniform actions): {}", reward_uniform);
    println!("  Reward (biased HOLD): {}", reward_biased);

    // Biased HOLD should have lower reward due to diversity penalty
    assert!(
        reward_biased < reward_uniform,
        "Diversity penalty not working: biased={}, uniform={}. Expected biased < uniform.",
        reward_biased,
        reward_uniform
    );

    Ok(())
}

#[tokio::test]
async fn test_batch_action_selection_consistency() -> Result<()> {
    // **TEST OBJECTIVE**: Verify that batched action selection produces same results
    // as sequential action selection (within randomness tolerance).
    //
    // **METHOD**: Select actions for same states in both modes, compare distributions.

    let hyperparams = create_test_hyperparams();
    let mut trainer = DQNTrainer::new(hyperparams)?;

    let training_data = create_synthetic_uptrend_data();

    // Extract 10 states
    let states: Result<Vec<_>> = training_data
        .iter()
        .take(10)
        .map(|(features, _)| {
            let close_price = Decimal::try_from(features[3]).unwrap_or(Decimal::ZERO);
            trainer.feature_vector_to_state(features, Some(close_price))
        })
        .collect();
    let states = states?;

    // Batched action selection
    let actions_batch = trainer.select_actions_batch(&states).await?;

    // Sequential action selection
    let mut actions_sequential = Vec::new();
    for state in &states {
        let action = trainer.select_action(state).await?;
        actions_sequential.push(action);
    }

    println!("Batch vs Sequential Action Selection:");
    println!("  Batch: {:?}", actions_batch);
    println!("  Sequential: {:?}", actions_sequential);

    // Count distributions (should be similar, but not identical due to epsilon-greedy randomness)
    let batch_hold_count = actions_batch
        .iter()
        .filter(|&&a| a == TradingAction::Hold)
        .count();
    let seq_hold_count = actions_sequential
        .iter()
        .filter(|&&a| a == TradingAction::Hold)
        .count();

    println!("  Batch HOLD count: {}", batch_hold_count);
    println!("  Sequential HOLD count: {}", seq_hold_count);

    // Both should have similar HOLD counts (within 20% tolerance)
    let diff = (batch_hold_count as i32 - seq_hold_count as i32).abs();
    assert!(
        diff <= 2,
        "Batch and sequential action selection differ significantly: batch={}, seq={}, diff={}",
        batch_hold_count,
        seq_hold_count,
        diff
    );

    Ok(())
}