foxhunt/crates/ml/tests/softmax_exploration_test.rs

#![allow(
    clippy::assertions_on_constants,
    clippy::assertions_on_result_states,
    clippy::clone_on_copy,
    clippy::decimal_literal_representation,
    clippy::doc_markdown,
    clippy::empty_line_after_doc_comments,
    clippy::field_reassign_with_default,
    clippy::get_unwrap,
    clippy::identity_op,
    clippy::inconsistent_digit_grouping,
    clippy::indexing_slicing,
    clippy::integer_division,
    clippy::len_zero,
    clippy::let_underscore_must_use,
    clippy::manual_div_ceil,
    clippy::manual_let_else,
    clippy::manual_range_contains,
    clippy::modulo_arithmetic,
    clippy::needless_range_loop,
    clippy::non_ascii_literal,
    clippy::redundant_clone,
    clippy::shadow_reuse,
    clippy::shadow_same,
    clippy::shadow_unrelated,
    clippy::single_match_else,
    clippy::str_to_string,
    clippy::string_slice,
    clippy::tests_outside_test_module,
    clippy::too_many_lines,
    clippy::unnecessary_wraps,
    clippy::unseparated_literal_suffix,
    clippy::use_debug,
    clippy::useless_vec,
    clippy::wildcard_enum_match_arm,
    clippy::else_if_without_else,
    clippy::expect_used,
    clippy::missing_const_for_fn,
    clippy::similar_names,
    clippy::type_complexity,
    clippy::collapsible_else_if,
    clippy::doc_lazy_continuation,
    clippy::items_after_test_module,
    clippy::map_clone,
    clippy::multiple_unsafe_ops_per_block,
    clippy::unwrap_or_default,
    clippy::assign_op_pattern,
    clippy::needless_borrow,
    clippy::println_empty_string,
    clippy::unnecessary_cast,
    clippy::used_underscore_binding,
    clippy::create_dir,
    clippy::implicit_saturating_sub,
    clippy::exit,
    clippy::expect_fun_call,
    clippy::too_many_arguments,
    clippy::unnecessary_map_or,
    clippy::unwrap_used,
    dead_code,
    unused_imports,
    unused_variables,
    clippy::cloned_ref_to_slice_refs,
    clippy::neg_multiply,
    clippy::while_let_loop,
    clippy::bool_assert_comparison,
    clippy::excessive_precision,
    clippy::trivially_copy_pass_by_ref,
    clippy::op_ref,
    clippy::redundant_closure,
    clippy::unnecessary_lazy_evaluations,
    clippy::if_then_some_else_none,
    clippy::unnecessary_to_owned,
    clippy::single_component_path_imports,
)]
//! Unit tests for softmax exploration implementation
//!
//! These tests verify the correctness of temperature-based softmax sampling
//! to replace argmax tie-breaking bias and fix action diversity collapse.
//!
//! Tests follow TDD methodology: written FIRST, before implementation.
//! Expected initial state: ALL TESTS SHOULD FAIL

use anyhow::Result;
use tracing::info;

/// Test 1: Softmax Distribution Basic
///
/// Verifies that softmax correctly computes probability distribution
/// where higher Q-values get higher (but not exclusive) probability.
#[test]
fn test_softmax_distribution_basic() -> Result<()> {
    // Q-values: [1.0, 2.0, 3.0]
    let q_values = [1.0_f32, 2.0, 3.0];
    let temperature = 1.0;

    // Call softmax function (host-side, takes &[f32])
    let probs_vec = ml::dqn::softmax::softmax_with_temperature(&q_values, temperature)?;

    // Action 2 (Q=3.0) should have highest probability (~66%)
    assert!(probs_vec[2] > probs_vec[1], "Highest Q-value should have highest probability");
    assert!(probs_vec[1] > probs_vec[0], "Middle Q-value should have middle probability");

    // All probabilities should be non-zero (exploration possible)
    assert!(probs_vec[0] > 0.0, "Lowest Q-value should still have non-zero probability");
    assert!(probs_vec[1] > 0.0, "Middle Q-value should have non-zero probability");
    assert!(probs_vec[2] > 0.0, "Highest Q-value should have non-zero probability");

    // Probabilities should sum to 1.0
    let sum: f32 = probs_vec.iter().sum();
    assert!((sum - 1.0).abs() < 1e-5, "Probabilities should sum to 1.0, got {}", sum);

    // Action 2 probability should be approximately 66% (softmax of [1, 2, 3])
    // Exact: exp(3) / (exp(1) + exp(2) + exp(3)) = 20.09 / 30.19 = 0.665
    assert!(
        (probs_vec[2] - 0.665).abs() < 0.01,
        "Action 2 probability should be ~66%, got {}",
        probs_vec[2]
    );

    Ok(())
}

/// Test 2: Temperature Effect on Distribution
///
/// Verifies that temperature controls exploration vs exploitation:
/// - High temp (10.0): Near-uniform distribution
/// - Low temp (0.1): Near-greedy (argmax-like)
#[test]
fn test_temperature_effect() -> Result<()> {
    let q_values = [1.0_f32, 2.0, 3.0];

    // High temperature: Should give near-uniform distribution
    let probs_high_vec = ml::dqn::softmax::softmax_with_temperature(&q_values, 10.0)?;

    // Calculate entropy for high temperature (should be high)
    let entropy_high: f32 = probs_high_vec
        .iter()
        .filter(|&&p| p > 0.0)
        .map(|&p| -p * p.log2())
        .sum();

    // Max entropy for 3 actions is log2(3) ~ 1.585
    assert!(
        entropy_high > 1.3,
        "High temperature should give high entropy (>1.3), got {}",
        entropy_high
    );

    // Low temperature: Should be near-greedy
    let probs_low_vec = ml::dqn::softmax::softmax_with_temperature(&q_values, 0.1)?;

    // Calculate entropy for low temperature (should be low)
    let entropy_low: f32 = probs_low_vec
        .iter()
        .filter(|&&p| p > 0.0)
        .map(|&p| -p * p.log2())
        .sum();

    // Low temperature should give low entropy (<0.5)
    assert!(
        entropy_low < 0.5,
        "Low temperature should give low entropy (<0.5), got {}",
        entropy_low
    );

    // Highest Q-value should dominate at low temperature (>99%)
    assert!(
        probs_low_vec[2] > 0.99,
        "Low temperature should make highest Q-value dominant (>99%), got {}",
        probs_low_vec[2]
    );

    Ok(())
}

/// Test 3: Tie-Breaking Fairness
///
/// When all Q-values are equal, softmax should give uniform distribution
/// (unlike argmax which always picks first action).
#[test]
fn test_tie_breaking_fairness() -> Result<()> {
    // All Q-values equal
    let q_values = [1.0_f32; 45];
    let temperature = 1.0;

    let probs_vec = ml::dqn::softmax::softmax_with_temperature(&q_values, temperature)?;

    // All probabilities should be equal (uniform distribution)
    let expected_prob = 1.0 / 45.0; // ~2.22%

    for (i, &prob) in probs_vec.iter().enumerate() {
        assert!(
            (prob - expected_prob).abs() < 1e-5,
            "Action {} probability should be ~{}, got {}",
            i,
            expected_prob,
            prob
        );
    }

    Ok(())
}

/// Test 4: Action Diversity Over Epochs (Sampling Test)
///
/// Verifies that repeated sampling from softmax produces diverse actions
/// (unlike argmax which always picks the same action).
#[test]
fn test_action_diversity_sampling() -> Result<()> {
    // Q-values with clear winner but other options viable
    let q_values = [1.0_f32, 1.5, 2.0];
    let temperature = 1.0;

    // Sample 1000 actions
    let num_samples = 1000;
    let mut action_counts = [0; 3];

    for _ in 0..num_samples {
        let action = ml::dqn::softmax::sample_from_softmax(&q_values, temperature)?;
        action_counts[action as usize] += 1;
    }

    // All actions should be selected at least once (diversity check)
    for (i, &count) in action_counts.iter().enumerate() {
        assert!(
            count > 0,
            "Action {} should be selected at least once in {} samples, got {}",
            i,
            num_samples,
            count
        );
    }

    // Diversity metric: percentage of actions used
    let diversity = action_counts.iter().filter(|&&c| c > 0).count() as f32 / 3.0;
    assert_eq!(
        diversity, 1.0,
        "Should use 100% of actions (3/3), got {:.1}%",
        diversity * 100.0
    );

    // Action 2 (highest Q-value) should be selected most often
    assert!(
        action_counts[2] > action_counts[1] && action_counts[1] > action_counts[0],
        "Higher Q-values should be selected more often: {:?}",
        action_counts
    );

    // But action 2 should NOT dominate completely (softmax property)
    let action2_ratio = action_counts[2] as f32 / num_samples as f32;
    assert!(
        action2_ratio < 0.8,
        "Highest Q-value should not dominate (>80%), got {:.1}%",
        action2_ratio * 100.0
    );

    Ok(())
}

/// Test 5: Numerical Stability with Extreme Q-values
///
/// Verifies that softmax handles extreme Q-values without NaN/Inf
/// (using log-sum-exp trick internally).
#[test]
fn test_numerical_stability_extreme_values() -> Result<()> {
    // Extreme Q-values that would overflow naive exp()
    let q_values = [-1000.0_f32, 0.0, 1000.0];
    let temperature = 1.0;

    // Should not panic or produce NaN/Inf
    let probs_vec = ml::dqn::softmax::softmax_with_temperature(&q_values, temperature)?;

    // Verify no NaN/Inf
    for (i, &prob) in probs_vec.iter().enumerate() {
        assert!(
            prob.is_finite(),
            "Action {} probability should be finite, got {}",
            i,
            prob
        );
        assert!(
            !prob.is_nan(),
            "Action {} probability should not be NaN",
            i
        );
    }

    // Probabilities should still sum to 1.0
    let sum: f32 = probs_vec.iter().sum();
    assert!(
        (sum - 1.0).abs() < 1e-5,
        "Probabilities should sum to 1.0 even with extreme values, got {}",
        sum
    );

    // Highest Q-value should dominate (but finite)
    assert!(
        probs_vec[2] > 0.99,
        "Extreme positive Q-value should dominate, got {}",
        probs_vec[2]
    );

    Ok(())
}

/// Test 6: Entropy Calculation
///
/// Verifies entropy metric for monitoring exploration level.
#[test]
fn test_softmax_entropy_calculation() -> Result<()> {
    // Test 1: Uniform distribution (max entropy)
    let q_uniform = [0.0_f32; 3];
    let entropy_uniform = ml::dqn::softmax::softmax_entropy(&q_uniform, 1.0)?;

    // Max entropy for 3 actions: log2(3) ~ 1.585
    assert!(
        (entropy_uniform - 1.585).abs() < 0.01,
        "Uniform distribution should have entropy ~1.585, got {}",
        entropy_uniform
    );

    // Test 2: Deterministic distribution (zero entropy)
    let q_deterministic = [-1000.0_f32, 0.0, 1000.0];
    let entropy_deterministic = ml::dqn::softmax::softmax_entropy(&q_deterministic, 0.1)?;

    // Near-deterministic should have very low entropy (<0.1)
    assert!(
        entropy_deterministic < 0.1,
        "Near-deterministic distribution should have low entropy, got {}",
        entropy_deterministic
    );

    Ok(())
}

/// Test 7: Batch Support
///
/// Verifies that softmax works with batched Q-values (multiple states).
#[test]
fn test_softmax_batch_support() -> Result<()> {
    // Batch of 4 states, 3 actions each
    let q_values_batch: Vec<Vec<f32>> = vec![
        vec![1.0_f32, 2.0, 3.0],
        vec![3.0, 2.0, 1.0],
        vec![2.0, 2.0, 2.0],
        vec![1.0, 1.5, 2.0],
    ];

    let temperature = 1.0;
    let probs_batch =
        ml::dqn::softmax::softmax_with_temperature_batch(&q_values_batch, temperature)?;

    // Verify shape
    assert_eq!(probs_batch.len(), 4, "Batch softmax should have 4 rows");
    for row in &probs_batch {
        assert_eq!(row.len(), 3, "Each row should have 3 action probabilities");
    }

    // Verify each row sums to 1.0
    for (i, row) in probs_batch.iter().enumerate() {
        let sum: f32 = row.iter().sum();
        assert!(
            (sum - 1.0).abs() < 1e-5,
            "Row {} probabilities should sum to 1.0, got {}",
            i,
            sum
        );
    }

    Ok(())
}

/// Test 8: 45-Action Space Support
///
/// Verifies softmax works correctly with full 45-action factored space.
#[test]
fn test_45_action_space_support() -> Result<()> {
    // Q-values for 45 actions (realistic scenario)
    let mut q_vec = vec![0.0_f32; 45];
    // Make some actions clearly better
    q_vec[10] = 2.0; // Good action
    q_vec[20] = 1.5; // Medium action
    q_vec[30] = 1.0; // Weak action

    let temperature = 1.0;

    let probs_vec = ml::dqn::softmax::softmax_with_temperature(&q_vec, temperature)?;

    // Verify 45 probabilities
    assert_eq!(probs_vec.len(), 45, "Should have 45 action probabilities");

    // Sum to 1.0
    let sum: f32 = probs_vec.iter().sum();
    assert!(
        (sum - 1.0).abs() < 1e-5,
        "45-action probabilities should sum to 1.0, got {}",
        sum
    );

    // Best actions should have higher probabilities
    assert!(
        probs_vec[10] > probs_vec[20] && probs_vec[20] > probs_vec[30],
        "Higher Q-values should have higher probabilities"
    );

    // But all actions should still have some probability (exploration)
    let num_nonzero = probs_vec.iter().filter(|&&p| p > 1e-6).count();
    assert!(
        num_nonzero >= 40,
        "At least 40/45 actions should have non-negligible probability, got {}",
        num_nonzero
    );

    Ok(())
}

/// Test 9: Chi-Squared Goodness of Fit
///
/// Statistical test to verify softmax sampling follows theoretical distribution.
#[test]
fn test_chi_squared_goodness_of_fit() -> Result<()> {
    // Simple 3-action case with known softmax probabilities
    let q_values = [0.0_f32, 1.0, 2.0];
    let temperature = 1.0;

    // Theoretical probabilities (computed via Python numpy):
    // exp(0)/sum = 1.0 / 11.1073 = 0.0900
    // exp(1)/sum = 2.718 / 11.1073 = 0.2447
    // exp(2)/sum = 7.389 / 11.1073 = 0.6652
    let expected = [0.0900, 0.2447, 0.6652];

    // Sample 10,000 times
    let num_samples = 10_000;
    let mut observed = [0; 3];

    for _ in 0..num_samples {
        let action = ml::dqn::softmax::sample_from_softmax(&q_values, temperature)?;
        observed[action as usize] += 1;
    }

    // Chi-squared test: sum((O_i - E_i)^2 / E_i)
    let mut chi_squared = 0.0;
    for i in 0..3 {
        let o = observed[i] as f32;
        let e = expected[i] * num_samples as f32;
        chi_squared += (o - e).powi(2) / e;
    }

    // Critical value for 2 degrees of freedom at p=0.05 is 5.99
    // We should NOT reject null hypothesis (sampling matches theoretical)
    assert!(
        chi_squared < 5.99,
        "Chi-squared test failed: statistic {} exceeds critical value 5.99 (samples don't match theoretical distribution)",
        chi_squared
    );

    info!(chi_squared, "Chi-squared statistic (critical value: 5.99)");

    Ok(())
}

/// Test 10: Zero Temperature Edge Case
///
/// Verifies that temperature=0.0 is handled gracefully (should behave like argmax).
#[test]
fn test_zero_temperature_edge_case() -> Result<()> {
    let q_values = [1.0_f32, 2.0, 3.0];

    // Temperature approaching zero should give near-deterministic distribution
    let probs_vec = ml::dqn::softmax::softmax_with_temperature(&q_values, 0.001)?;

    // Highest Q-value should have ~100% probability
    assert!(
        probs_vec[2] > 0.9999,
        "Near-zero temperature should make highest Q-value dominant (>99.99%), got {}",
        probs_vec[2]
    );

    // Other actions should have negligible probability
    assert!(
        probs_vec[0] < 0.0001 && probs_vec[1] < 0.0001,
        "Near-zero temperature should make other actions negligible, got [{}, {}]",
        probs_vec[0],
        probs_vec[1]
    );

    Ok(())
}