foxhunt/ml/tests/dqn_ensemble_tests.rs

//! DQN Ensemble Oracle Integration Tests
//!
//! Comprehensive end-to-end tests validating the ensemble oracle system with multi-model voting:
//! 1. Voting Mechanisms (majority voting, tie-breaking)
//! 2. Uncertainty Metrics (diversity bonus, agreement calculation)
//! 3. Training Convergence (ensemble vs single-agent comparison)
//! 4. Edge Cases (empty votes, single model, all disagree)
//!
//! Test Categories:
//! - Voting Mechanisms (6 tests)
//! - Uncertainty Metrics (5 tests)
//! - Graceful Degradation (4 tests)
//! - Integration Tests (3 tests)
//! - Training Comparison (2 tests)
//!
//! Total: 20 tests

#![allow(unused_crate_dependencies)]

use anyhow::Result;
use candle_core::{Device, Tensor};
use ml::dqn::{agent::TradingAction, ensemble_oracle::EnsembleOracle};
use ml::MLError;

// ================================================================================================
// TEST UTILITIES MODULE
// ================================================================================================

mod test_utils {
    use super::*;

    /// Test fixture for Ensemble Oracle tests
    pub struct EnsembleTestFixture {
        pub device: Device,
        pub oracle_enabled: EnsembleOracle,
        pub oracle_disabled: EnsembleOracle,
        pub mock_state: Tensor,
    }

    impl EnsembleTestFixture {
        pub fn new() -> Result<Self, MLError> {
            let device = Device::Cpu;

            // Create enabled oracle (bypasses model loading stub)
            let mut oracle_enabled = EnsembleOracle::new();
            oracle_enabled.load_models(Some("path/to/transformer"), None, None)?;

            // Create disabled oracle (no models loaded)
            let oracle_disabled = EnsembleOracle::new();

            // Mock state tensor (128-dimensional)
            let mock_state = Tensor::randn(0.0f32, 1.0f32, &[1, 128], &device)?;

            Ok(Self {
                device,
                oracle_enabled,
                oracle_disabled,
                mock_state,
            })
        }

        /// Helper: Assert reward is within valid range [0.0, 0.8]
        pub fn assert_reward_valid(reward: f64) {
            assert!(reward.is_finite(), "Reward must be finite, got: {}", reward);
            assert!(
                reward >= 0.0,
                "Reward must be non-negative, got: {}",
                reward
            );
            assert!(reward <= 0.8, "Reward must be ≤ 0.8, got: {}", reward);
        }

        /// Helper: Calculate reward with votes
        pub fn calculate_reward(
            &self,
            oracle: &EnsembleOracle,
            action: TradingAction,
            votes: Vec<usize>,
        ) -> f64 {
            oracle.calculate_ensemble_reward(&self.mock_state, action, votes)
        }
    }
}

// ================================================================================================
// CATEGORY 1: VOTING MECHANISMS (6 TESTS)
// ================================================================================================

#[test]
fn test_majority_voting_clear_winner() -> Result<()> {
    let fixture = test_utils::EnsembleTestFixture::new()?;

    // 2 models vote BUY, 1 model votes SELL → BUY is majority
    let votes = vec![0, 0, 1]; // BUY, BUY, SELL
    let reward = fixture.calculate_reward(&fixture.oracle_enabled, TradingAction::Buy, votes);

    // DQN agrees with majority → agreement=0.5, diversity=0.1 (2 unique) = 0.6
    assert!(
        (reward - 0.6).abs() < 1e-6,
        "Expected 0.6 (0.5 + 0.1), got {}",
        reward
    );
    test_utils::EnsembleTestFixture::assert_reward_valid(reward);

    Ok(())
}

#[test]
fn test_majority_voting_disagreement() -> Result<()> {
    let fixture = test_utils::EnsembleTestFixture::new()?;

    // 2 models vote SELL, 1 model votes BUY → SELL is majority
    let votes = vec![1, 1, 0]; // SELL, SELL, BUY
    let reward = fixture.calculate_reward(&fixture.oracle_enabled, TradingAction::Buy, votes);

    // DQN disagrees with majority → agreement=0.1, diversity=0.1 (2 unique) = 0.2
    assert!(
        (reward - 0.2).abs() < 1e-6,
        "Expected 0.2 (0.1 + 0.1), got {}",
        reward
    );
    test_utils::EnsembleTestFixture::assert_reward_valid(reward);

    Ok(())
}

#[test]
fn test_majority_voting_three_way_tie() -> Result<()> {
    let fixture = test_utils::EnsembleTestFixture::new()?;

    // All 3 models vote different actions → tie-breaking is non-deterministic
    let votes = vec![0, 1, 2]; // BUY, SELL, HOLD
    let reward = fixture.calculate_reward(&fixture.oracle_enabled, TradingAction::Buy, votes);

    // With 3-way tie, majority is picked by HashMap iteration order (non-deterministic)
    // Reward can be 0.8 (if BUY is majority) or 0.4 (if SELL/HOLD is majority)
    // Both are valid, diversity bonus is 0.3 for 3 unique actions
    assert!(
        (reward - 0.8).abs() < 1e-6 || (reward - 0.4).abs() < 1e-6,
        "Expected 0.8 or 0.4, got {}",
        reward
    );
    test_utils::EnsembleTestFixture::assert_reward_valid(reward);

    Ok(())
}

#[test]
fn test_majority_voting_full_consensus() -> Result<()> {
    let fixture = test_utils::EnsembleTestFixture::new()?;

    // All models agree on BUY → full consensus
    let votes = vec![0, 0, 0]; // BUY, BUY, BUY
    let reward = fixture.calculate_reward(&fixture.oracle_enabled, TradingAction::Buy, votes);

    // DQN agrees with unanimous consensus → agreement=0.5, diversity=0.0 (1 unique) = 0.5
    assert!(
        (reward - 0.5).abs() < 1e-6,
        "Expected 0.5 (0.5 + 0.0), got {}",
        reward
    );
    test_utils::EnsembleTestFixture::assert_reward_valid(reward);

    Ok(())
}

#[test]
fn test_majority_voting_two_vs_one() -> Result<()> {
    let fixture = test_utils::EnsembleTestFixture::new()?;

    // 2 BUY, 1 SELL → BUY is clear majority
    let votes = vec![0, 1, 0]; // BUY, SELL, BUY
    let reward = fixture.calculate_reward(&fixture.oracle_enabled, TradingAction::Buy, votes);

    // DQN agrees with majority (2 BUY) → agreement=0.5, diversity=0.1 (2 unique) = 0.6
    assert!((reward - 0.6).abs() < 1e-6, "Expected 0.6, got {}", reward);
    test_utils::EnsembleTestFixture::assert_reward_valid(reward);

    Ok(())
}

#[test]
fn test_voting_with_hold_action() -> Result<()> {
    let fixture = test_utils::EnsembleTestFixture::new()?;

    // 2 models vote HOLD, 1 votes BUY → HOLD is majority
    let votes = vec![2, 2, 0]; // HOLD, HOLD, BUY
    let reward = fixture.calculate_reward(&fixture.oracle_enabled, TradingAction::Hold, votes);

    // DQN agrees with HOLD majority → agreement=0.5, diversity=0.1 (2 unique) = 0.6
    assert!((reward - 0.6).abs() < 1e-6, "Expected 0.6, got {}", reward);
    test_utils::EnsembleTestFixture::assert_reward_valid(reward);

    Ok(())
}

// ================================================================================================
// CATEGORY 2: UNCERTAINTY METRICS (5 TESTS)
// ================================================================================================

#[test]
fn test_diversity_bonus_high_uncertainty() -> Result<()> {
    let fixture = test_utils::EnsembleTestFixture::new()?;

    // All models disagree → high uncertainty, diversity bonus = 0.3
    let votes = vec![0, 1, 2]; // BUY, SELL, HOLD
    let reward = fixture.calculate_reward(&fixture.oracle_enabled, TradingAction::Buy, votes);

    // Diversity bonus is 0.3 (3 unique actions), agreement bonus is 0.1 or 0.5 (non-deterministic tie)
    // Total reward: 0.4 or 0.8
    assert!(
        reward >= 0.4 && reward <= 0.8,
        "Expected reward in [0.4, 0.8], got {}",
        reward
    );
    test_utils::EnsembleTestFixture::assert_reward_valid(reward);

    Ok(())
}

#[test]
fn test_diversity_bonus_moderate_uncertainty() -> Result<()> {
    let fixture = test_utils::EnsembleTestFixture::new()?;

    // 2 unique actions → moderate uncertainty, diversity bonus = 0.1
    let votes = vec![0, 0, 1]; // BUY, BUY, SELL
    let reward = fixture.calculate_reward(&fixture.oracle_enabled, TradingAction::Buy, votes);

    // Agreement=0.5 (DQN agrees with BUY majority), diversity=0.1 (2 unique) = 0.6
    assert!((reward - 0.6).abs() < 1e-6, "Expected 0.6, got {}", reward);
    test_utils::EnsembleTestFixture::assert_reward_valid(reward);

    Ok(())
}

#[test]
fn test_diversity_bonus_no_uncertainty() -> Result<()> {
    let fixture = test_utils::EnsembleTestFixture::new()?;

    // All models agree → no uncertainty, diversity bonus = 0.0
    let votes = vec![0, 0, 0]; // BUY, BUY, BUY
    let reward = fixture.calculate_reward(&fixture.oracle_enabled, TradingAction::Buy, votes);

    // Agreement=0.5 (DQN agrees), diversity=0.0 (1 unique) = 0.5
    assert!((reward - 0.5).abs() < 1e-6, "Expected 0.5, got {}", reward);
    test_utils::EnsembleTestFixture::assert_reward_valid(reward);

    Ok(())
}

#[test]
fn test_agreement_bonus_high() -> Result<()> {
    let fixture = test_utils::EnsembleTestFixture::new()?;

    // DQN agrees with majority → high agreement bonus = 0.5
    let votes = vec![0, 0, 1]; // BUY, BUY, SELL
    let reward = fixture.calculate_reward(&fixture.oracle_enabled, TradingAction::Buy, votes);

    // Agreement=0.5 (DQN agrees with BUY majority), diversity=0.1 = 0.6
    assert!(reward >= 0.5, "Expected reward ≥ 0.5, got {}", reward);
    test_utils::EnsembleTestFixture::assert_reward_valid(reward);

    Ok(())
}

#[test]
fn test_agreement_bonus_low() -> Result<()> {
    let fixture = test_utils::EnsembleTestFixture::new()?;

    // DQN disagrees with majority → low agreement bonus = 0.1
    let votes = vec![1, 1, 0]; // SELL, SELL, BUY
    let reward = fixture.calculate_reward(&fixture.oracle_enabled, TradingAction::Buy, votes);

    // Agreement=0.1 (DQN disagrees with SELL majority), diversity=0.1 = 0.2
    assert!(reward <= 0.3, "Expected reward ≤ 0.3, got {}", reward);
    test_utils::EnsembleTestFixture::assert_reward_valid(reward);

    Ok(())
}

// ================================================================================================
// CATEGORY 3: GRACEFUL DEGRADATION (4 TESTS)
// ================================================================================================

#[test]
fn test_disabled_oracle_returns_zero() -> Result<()> {
    let fixture = test_utils::EnsembleTestFixture::new()?;

    // Disabled oracle (no models loaded) → should return 0.0 gracefully
    let votes = vec![0, 1, 2];
    let reward = fixture.calculate_reward(&fixture.oracle_disabled, TradingAction::Buy, votes);

    assert!((reward - 0.0).abs() < 1e-6, "Expected 0.0, got {}", reward);

    Ok(())
}

#[test]
fn test_empty_votes_returns_zero() -> Result<()> {
    let fixture = test_utils::EnsembleTestFixture::new()?;

    // Empty votes → should return 0.0 gracefully (edge case)
    let votes = vec![];
    let reward = fixture.calculate_reward(&fixture.oracle_enabled, TradingAction::Buy, votes);

    assert!((reward - 0.0).abs() < 1e-6, "Expected 0.0, got {}", reward);

    Ok(())
}

#[test]
fn test_single_model_no_diversity() -> Result<()> {
    let fixture = test_utils::EnsembleTestFixture::new()?;

    // Only 1 model loaded → no diversity, but agreement still counts
    let votes = vec![1]; // SELL
    let reward = fixture.calculate_reward(&fixture.oracle_enabled, TradingAction::Sell, votes);

    // Agreement=0.5 (DQN agrees with single SELL vote), diversity=0.0 (1 unique) = 0.5
    assert!((reward - 0.5).abs() < 1e-6, "Expected 0.5, got {}", reward);
    test_utils::EnsembleTestFixture::assert_reward_valid(reward);

    Ok(())
}

#[test]
fn test_single_model_disagreement() -> Result<()> {
    let fixture = test_utils::EnsembleTestFixture::new()?;

    // Only 1 model, DQN disagrees → low agreement bonus
    let votes = vec![1]; // SELL
    let reward = fixture.calculate_reward(&fixture.oracle_enabled, TradingAction::Buy, votes);

    // Agreement=0.1 (DQN disagrees with SELL), diversity=0.0 (1 unique) = 0.1
    assert!((reward - 0.1).abs() < 1e-6, "Expected 0.1, got {}", reward);
    test_utils::EnsembleTestFixture::assert_reward_valid(reward);

    Ok(())
}

// ================================================================================================
// CATEGORY 4: INTEGRATION TESTS (3 TESTS)
// ================================================================================================

#[test]
fn test_model_loading_enables_oracle() -> Result<()> {
    let mut oracle = EnsembleOracle::new();
    assert!(
        !oracle.load_models(None, None, None).is_err(),
        "load_models should not error"
    );

    // Load 1 model
    oracle.load_models(Some("path/to/transformer"), None, None)?;

    // Oracle should be enabled
    let votes = vec![0, 1, 2];
    let mock_state = Tensor::randn(0.0f32, 1.0f32, &[1, 128], &Device::Cpu)?;
    let reward = oracle.calculate_ensemble_reward(&mock_state, TradingAction::Buy, votes);

    // Enabled oracle should return non-zero reward
    assert!(reward > 0.0, "Expected non-zero reward, got {}", reward);

    Ok(())
}

#[test]
fn test_model_loading_multiple_models() -> Result<()> {
    let mut oracle = EnsembleOracle::new();

    // Load multiple models
    oracle.load_models(
        Some("path/to/transformer"),
        Some("path/to/lstm"),
        Some("path/to/ppo"),
    )?;

    // Oracle should be enabled
    let votes = vec![0, 1, 2];
    let mock_state = Tensor::randn(0.0f32, 1.0f32, &[1, 128], &Device::Cpu)?;
    let reward = oracle.calculate_ensemble_reward(&mock_state, TradingAction::Buy, votes);

    // Enabled oracle should return non-zero reward
    assert!(reward > 0.0, "Expected non-zero reward, got {}", reward);

    Ok(())
}

#[test]
fn test_model_loading_no_models_disables_oracle() -> Result<()> {
    let mut oracle = EnsembleOracle::new();

    // Load no models
    oracle.load_models(None, None, None)?;

    // Oracle should remain disabled
    let votes = vec![0, 1, 2];
    let mock_state = Tensor::randn(0.0f32, 1.0f32, &[1, 128], &Device::Cpu)?;
    let reward = oracle.calculate_ensemble_reward(&mock_state, TradingAction::Buy, votes);

    // Disabled oracle should return 0.0
    assert!((reward - 0.0).abs() < 1e-6, "Expected 0.0, got {}", reward);

    Ok(())
}

// ================================================================================================
// CATEGORY 5: TRAINING COMPARISON (2 TESTS)
// ================================================================================================

/// Smoke test: Ensemble vs single-agent reward comparison over 5 epochs
///
/// This test validates that ensemble rewards are within expected bounds
/// compared to single-agent training. It does NOT run full training,
/// but simulates reward calculations for 5 mock episodes.
#[test]
fn test_ensemble_vs_single_agent_rewards() -> Result<()> {
    let fixture = test_utils::EnsembleTestFixture::new()?;

    // Simulate 5 episodes with different vote patterns
    let test_cases = vec![
        (vec![0, 0, 0], TradingAction::Buy, "Full consensus"),
        (vec![0, 0, 1], TradingAction::Buy, "Majority agreement"),
        (vec![0, 1, 2], TradingAction::Buy, "High disagreement"),
        (vec![1, 1, 0], TradingAction::Buy, "Majority disagreement"),
        (vec![2, 2, 2], TradingAction::Hold, "HOLD consensus"),
    ];

    for (votes, action, scenario) in test_cases {
        let reward = fixture.calculate_reward(&fixture.oracle_enabled, action, votes);

        // Validate reward bounds for all scenarios
        test_utils::EnsembleTestFixture::assert_reward_valid(reward);
        println!("Scenario '{}': reward={:.4}", scenario, reward);

        // Ensemble rewards should be in [0.0, 0.8]
        assert!(
            reward <= 0.8,
            "Reward exceeds max 0.8 for scenario '{}'",
            scenario
        );
    }

    Ok(())
}

/// Smoke test: Convergence pattern validation
///
/// Validates that ensemble rewards show expected patterns:
/// - High diversity (3 unique votes) → reward in [0.4, 0.8]
/// - Moderate diversity (2 unique votes) → reward in [0.2, 0.6]
/// - No diversity (1 unique vote) → reward in [0.1, 0.5]
#[test]
fn test_ensemble_convergence_patterns() -> Result<()> {
    let fixture = test_utils::EnsembleTestFixture::new()?;

    // Test diversity patterns
    let high_diversity = vec![0, 1, 2]; // 3 unique actions
    let moderate_diversity = vec![0, 0, 1]; // 2 unique actions
    let no_diversity = vec![0, 0, 0]; // 1 unique action

    let reward_high =
        fixture.calculate_reward(&fixture.oracle_enabled, TradingAction::Buy, high_diversity);
    let reward_moderate = fixture.calculate_reward(
        &fixture.oracle_enabled,
        TradingAction::Buy,
        moderate_diversity,
    );
    let reward_no =
        fixture.calculate_reward(&fixture.oracle_enabled, TradingAction::Buy, no_diversity);

    // Validate diversity patterns
    println!("High diversity reward: {:.4}", reward_high);
    println!("Moderate diversity reward: {:.4}", reward_moderate);
    println!("No diversity reward: {:.4}", reward_no);

    // High diversity should have highest potential reward
    assert!(
        reward_high >= 0.4,
        "High diversity reward too low: {}",
        reward_high
    );
    assert!(
        reward_high <= 0.8,
        "High diversity reward too high: {}",
        reward_high
    );

    // Moderate diversity should be mid-range
    assert!(
        reward_moderate >= 0.2,
        "Moderate diversity reward too low: {}",
        reward_moderate
    );
    assert!(
        reward_moderate <= 0.6,
        "Moderate diversity reward too high: {}",
        reward_moderate
    );

    // No diversity should have lowest reward (but still valid)
    assert!(
        reward_no >= 0.1,
        "No diversity reward too low: {}",
        reward_no
    );
    assert!(
        reward_no <= 0.5,
        "No diversity reward too high: {}",
        reward_no
    );

    Ok(())
}