foxhunt/ml/tests/epsilon_greedy_softmax_test.rs

// Test file for DQN softmax action selection
// Validates WorkingDQN::select_action behavior with epsilon-greedy + softmax

use ml::dqn::dqn::{WorkingDQN, WorkingDQNConfig};
use std::collections::HashMap;

/// Helper: Create DQN with custom epsilon/temperature
fn create_dqn_with_params(epsilon: f64, temperature: f64, state_dim: usize) -> WorkingDQN {
    // For test flexibility: allow temperature_min to be lower than temperature_start
    // This enables tests with very low temperatures (e.g., 0.01 for deterministic tests)
    let temperature_min = (temperature * 0.5).max(0.01); // Set min to half of start, floor at 0.01

    let mut config = WorkingDQNConfig {
        state_dim,
        num_actions: 3,
        hidden_dims: vec![64, 32],
        learning_rate: 0.001,
        gamma: 0.99,
        epsilon_start: epsilon as f32,
        epsilon_end: 0.01,
        epsilon_decay: 0.995,
        replay_buffer_capacity: 1000,
        batch_size: 32,
        min_replay_size: 32,
        target_update_freq: 100,
        use_double_dqn: false,
        use_huber_loss: true,
        huber_delta: 1.0,
        leaky_relu_alpha: 0.01,
        gradient_clip_norm: 10.0,
        tau: 0.001,
        use_soft_updates: false,
        warmup_steps: 0, // Disable warmup for testing
        temperature_start: temperature,
        temperature_min, // Adaptive based on temperature_start
        temperature_decay: 0.995,
        target_temperature_fraction: 0.75,
        variance_multiplier: 0.5,
        use_adaptive_temperature: false,
        loss_improvement_threshold: 0.999,
        plateau_window: 10,
        temp_increase_factor: 1.05,
        temperature_slow_decay: 0.998,
    };

    WorkingDQN::new(config).expect("Failed to create DQN")
}

/// Helper: Compute softmax probabilities manually
fn compute_softmax(q_values: &[f64], temperature: f64) -> Vec<f64> {
    let scaled: Vec<f64> = q_values.iter().map(|q| q / temperature).collect();
    let max_val = scaled.iter().cloned().fold(f64::NEG_INFINITY, f64::max);
    let exp_vals: Vec<f64> = scaled.iter().map(|x| (x - max_val).exp()).collect();
    let sum: f64 = exp_vals.iter().sum();
    exp_vals.iter().map(|x| x / sum).collect()
}

#[test]
fn test_softmax_equal_q_values() {
    println!("\n=== Test 1: Softmax with Equal Q-Values ===");

    let state_dim = 128;
    let mut dqn = create_dqn_with_params(0.0, 1.0, state_dim); // No exploration, temp=1.0
    let state = vec![0.0; state_dim]; // Zero state (will get random Q-values from untrained network)

    // Mock Q-values for reference
    let q_values = vec![100.0, 100.0, 100.0];
    let expected_probs = compute_softmax(&q_values, 1.0);

    println!("Q-values (reference): {:?}", q_values);
    println!("Expected probabilities: {:?}", expected_probs);

    // Sample 1000 actions
    let mut action_counts = HashMap::new();
    for _ in 0..1000 {
        let action = dqn.select_action(&state).unwrap();
        *action_counts.entry(action as u8).or_insert(0) += 1;
    }

    let observed_probs: Vec<f64> = (0..3)
        .map(|i| *action_counts.get(&i).unwrap_or(&0) as f64 / 1000.0)
        .collect();

    println!("Observed probabilities: {:?}", observed_probs);
    println!("Action counts: {:?}", action_counts);

    // With untrained network, Q-values will be near-random, so distribution should be somewhat uniform
    // We don't enforce strict uniformity since the network is random, just check all actions are taken
    for (i, &prob) in observed_probs.iter().enumerate() {
        assert!(
            prob > 0.0,
            "Action {} should be selected at least once, got probability {}",
            i,
            prob
        );
    }

    println!("✅ Test 1 PASSED: All actions selected with untrained network\n");
}

#[test]
fn test_action_frequency_alignment() {
    println!("\n=== Test 3: Action Frequency Alignment ===");

    let state_dim = 128;
    let mut dqn = create_dqn_with_params(0.0, 0.1, state_dim); // No exploration, low temp
    let state = vec![0.5; state_dim]; // Non-zero state

    println!("Sampling 1000 actions with epsilon=0.0, temperature=0.1");

    // Sample 1000 actions
    let mut action_counts = HashMap::new();
    for _ in 0..1000 {
        let action = dqn.select_action(&state).unwrap();
        *action_counts.entry(action as u8).or_insert(0) += 1;
    }

    let observed_probs: Vec<f64> = (0..3)
        .map(|i| *action_counts.get(&i).unwrap_or(&0) as f64 / 1000.0)
        .collect();

    println!("Observed probabilities: {:?}", observed_probs);
    println!("Action counts: {:?}", action_counts);

    // With untrained network + low temperature, one action should be most frequent
    // but may not dominate as strongly as with a trained network
    let max_prob = observed_probs.iter().cloned().fold(0.0, f64::max);
    assert!(
        max_prob > 0.35,
        "With low temperature (0.1), one action should be most frequent (>35%), got max={:.3}",
        max_prob
    );

    println!(
        "✅ Test 3 PASSED: Action selection working (max prob: {:.3})\n",
        max_prob
    );
    println!("NOTE: Untrained network Q-values are pseudo-random, so distribution varies.\n");
}

#[test]
fn test_inverse_q_value_check() {
    println!("\n=== Test 4: INVERSE CHECK (Critical Bug Detection) ===");
    println!("NOTE: This test uses an untrained network, so Q-values are pseudo-random.");
    println!("The critical check is whether action selection is consistent with Q-values.\n");

    let state_dim = 128;
    let mut dqn = create_dqn_with_params(0.0, 0.1, state_dim); // No exploration, low temp

    // Use different states to get different Q-value patterns
    let state1 = vec![1.0; state_dim]; // High positive state
    let state2 = vec![-1.0; state_dim]; // High negative state
    let state3 = vec![0.0; state_dim]; // Zero state

    println!("Testing 3 different states to verify action selection consistency:");

    for (idx, state) in [state1, state2, state3].iter().enumerate() {
        // Sample 1000 actions
        let mut action_counts = HashMap::new();
        for _ in 0..1000 {
            let action = dqn.select_action(state).unwrap();
            *action_counts.entry(action as u8).or_insert(0) += 1;
        }

        let observed_probs: Vec<f64> = (0..3)
            .map(|i| *action_counts.get(&i).unwrap_or(&0) as f64 / 1000.0)
            .collect();

        // Find most frequent action
        let (max_action, max_prob) = observed_probs
            .iter()
            .enumerate()
            .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap())
            .map(|(idx, prob)| (idx, *prob))
            .unwrap();

        println!(
            "State {}: Most frequent action {} (probability: {:.3})",
            idx + 1,
            max_action,
            max_prob
        );
        println!(
            "  Action distribution: BUY={:.3}, SELL={:.3}, HOLD={:.3}",
            observed_probs[0], observed_probs[1], observed_probs[2]
        );

        // With low temperature (0.1), one action should be most probable
        // Note: Untrained networks may have similar Q-values, so we use >40% threshold
        // (More lenient than 50% to account for random initialization variance)
        assert!(
            max_prob > 0.40,
            "State {}: Expected one action to be most probable (>40%), got max={:.3}",
            idx + 1,
            max_prob
        );
    }

    println!("\n✅ Test 4 PASSED: Action selection shows strong preference (low temp effect)\n");
}

#[test]
fn test_epsilon_greedy_exploration_mix() {
    println!("\n=== Test 5: Epsilon-Greedy Exploration Mix ===");

    let state_dim = 128;
    let mut dqn = create_dqn_with_params(0.3, 0.1, state_dim); // 30% exploration, temp=0.1
    let state = vec![0.5; state_dim];

    println!("Epsilon: 0.3 (30% random exploration)");
    println!("Temperature: 0.1 (amplifies differences when not exploring)");

    // Sample 1000 actions
    let mut action_counts = HashMap::new();
    for _ in 0..1000 {
        let action = dqn.select_action(&state).unwrap();
        *action_counts.entry(action as u8).or_insert(0) += 1;
    }

    let observed_probs: Vec<f64> = (0..3)
        .map(|i| *action_counts.get(&i).unwrap_or(&0) as f64 / 1000.0)
        .collect();

    println!("Observed probabilities: {:?}", observed_probs);
    println!("Action counts: {:?}", action_counts);

    // Find most frequent action
    let max_prob = observed_probs.iter().cloned().fold(0.0, f64::max);

    // With 30% epsilon + low temp, expect:
    // - One action should still be most frequent (from greedy 70% + some exploration)
    // - But not as extreme as pure greedy (should be < 90%)
    assert!(
        max_prob > 0.4 && max_prob < 0.9,
        "Expected max probability in range [0.4, 0.9] with epsilon=0.3, got {:.3}",
        max_prob
    );

    // All actions should be selected at least sometimes (due to exploration)
    for (i, &prob) in observed_probs.iter().enumerate() {
        assert!(
            prob > 0.05,
            "Action {} should be selected >5% due to exploration, got {:.3}",
            i,
            prob
        );
    }

    println!(
        "✅ Test 5 PASSED: Epsilon exploration adds randomness (max prob: {:.3})\n",
        max_prob
    );
}

#[test]
fn test_temperature_effect() {
    println!("\n=== Test 6: Temperature Effect on Action Distribution ===");

    let state_dim = 128;
    let state = vec![0.7; state_dim]; // Fixed state

    // Test three temperature values
    let temperatures = vec![0.1, 1.0, 10.0];

    for &temp in &temperatures {
        let mut dqn = create_dqn_with_params(0.0, temp, state_dim); // No exploration

        // Sample 1000 actions
        let mut action_counts = HashMap::new();
        for _ in 0..1000 {
            let action = dqn.select_action(&state).unwrap();
            *action_counts.entry(action as u8).or_insert(0) += 1;
        }

        let observed_probs: Vec<f64> = (0..3)
            .map(|i| *action_counts.get(&i).unwrap_or(&0) as f64 / 1000.0)
            .collect();

        let max_prob = observed_probs.iter().cloned().fold(0.0, f64::max);
        let min_prob = observed_probs.iter().cloned().fold(1.0, f64::min);

        println!(
            "Temperature {:.1}: max_prob={:.3}, min_prob={:.3}",
            temp, max_prob, min_prob
        );
        println!(
            "  Distribution: BUY={:.3}, SELL={:.3}, HOLD={:.3}",
            observed_probs[0], observed_probs[1], observed_probs[2]
        );

        // Low temperature (0.1): Strong preference (max > 0.6)
        if temp < 0.5 {
            assert!(
                max_prob > 0.6,
                "Low temperature ({}) should create strong preference, got max={:.3}",
                temp,
                max_prob
            );
        }

        // High temperature (10.0): More uniform distribution (max < 0.7)
        if temp > 5.0 {
            assert!(
                max_prob < 0.7,
                "High temperature ({}) should create more uniform distribution, got max={:.3}",
                temp,
                max_prob
            );
        }
    }

    println!("\n✅ Test 6 PASSED: Temperature correctly controls exploration/exploitation\n");
}

#[test]
fn test_deterministic_evaluation_mode() {
    println!("\n=== Test 7: Deterministic Evaluation (epsilon=0.0) ===");

    let state_dim = 128;
    let mut dqn = create_dqn_with_params(0.0, 0.01, state_dim); // Zero exploration, very low temp
    let state = vec![0.5; state_dim];

    println!("Epsilon: 0.0 (pure greedy)");
    println!("Temperature: 0.01 (near-deterministic softmax)");

    // Sample 100 actions
    let mut action_counts = HashMap::new();
    for _ in 0..100 {
        let action = dqn.select_action(&state).unwrap();
        *action_counts.entry(action as u8).or_insert(0) += 1;
    }

    let observed_probs: Vec<f64> = (0..3)
        .map(|i| *action_counts.get(&i).unwrap_or(&0) as f64 / 100.0)
        .collect();

    println!("Observed probabilities: {:?}", observed_probs);

    // With epsilon=0.0 and very low temp, one action should strongly dominate
    // Untrained network may not reach 100% determinism, but should be >90%
    let max_prob = observed_probs.iter().cloned().fold(0.0, f64::max);
    assert!(
        max_prob > 0.90,
        "Deterministic mode should select one action >90% of time, got {:.3}",
        max_prob
    );

    println!(
        "✅ Test 7 PASSED: Near-deterministic evaluation works (max prob: {:.3})\n",
        max_prob
    );
    println!("NOTE: Untrained network Q-values are pseudo-random.\n");
}