foxhunt/ml/tests/softmax_exploration_test.rs

//! 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 candle_core::{Device, Tensor};

/// 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<()> {
    let device = Device::Cpu;

    // Q-values: [1.0, 2.0, 3.0]
    let q_values = Tensor::new(&[1.0f32, 2.0, 3.0], &device)?;
    let temperature = 1.0;

    // Call softmax function (doesn't exist yet - will fail)
    let probs = ml::dqn::softmax::softmax_with_temperature(&q_values, temperature)?;

    // Extract probabilities
    let probs_vec = probs.to_vec1::<f32>()?;

    // 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 device = Device::Cpu;
    let q_values = Tensor::new(&[1.0f32, 2.0, 3.0], &device)?;

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

    // 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 = ml::dqn::softmax::softmax_with_temperature(&q_values, 0.1)?;
    let probs_low_vec = probs_low.to_vec1::<f32>()?;

    // 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<()> {
    let device = Device::Cpu;

    // All Q-values equal
    let q_values = Tensor::new(&[1.0f32; 45], &device)?;
    let temperature = 1.0;

    let probs = ml::dqn::softmax::softmax_with_temperature(&q_values, temperature)?;
    let probs_vec = probs.to_vec1::<f32>()?;

    // 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<()> {
    let device = Device::Cpu;

    // Q-values with clear winner but other options viable
    let q_values = Tensor::new(&[1.0f32, 1.5, 2.0], &device)?;
    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<()> {
    let device = Device::Cpu;

    // Extreme Q-values that would overflow naive exp()
    let q_values = Tensor::new(&[-1000.0f32, 0.0, 1000.0], &device)?;
    let temperature = 1.0;

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

    // 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<()> {
    let device = Device::Cpu;

    // Test 1: Uniform distribution (max entropy)
    let q_uniform = Tensor::new(&[0.0f32; 3], &device)?;
    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 = Tensor::new(&[-1000.0f32, 0.0, 1000.0], &device)?;
    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<()> {
    let device = Device::Cpu;

    // Batch of 4 states, 3 actions each
    let q_values = Tensor::new(
        &[
            [1.0f32, 2.0, 3.0],
            [3.0, 2.0, 1.0],
            [2.0, 2.0, 2.0],
            [1.0, 1.5, 2.0],
        ],
        &device,
    )?;

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

    // Verify shape
    assert_eq!(probs.dims(), &[4, 3], "Batch softmax should preserve shape");

    // Verify each row sums to 1.0
    let probs_2d = probs.to_vec2::<f32>()?;
    for (i, row) in probs_2d.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<()> {
    let device = Device::Cpu;

    // Q-values for 45 actions (realistic scenario)
    let mut q_vec = vec![0.0f32; 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 q_values = Tensor::new(q_vec.as_slice(), &device)?;
    let temperature = 1.0;

    let probs = ml::dqn::softmax::softmax_with_temperature(&q_values, temperature)?;
    let probs_vec = probs.to_vec1::<f32>()?;

    // 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<()> {
    let device = Device::Cpu;

    // Simple 3-action case with known softmax probabilities
    let q_values = Tensor::new(&[0.0f32, 1.0, 2.0], &device)?;
    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: Σ((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
    );

    println!("Chi-squared statistic: {:.2} (critical value: 5.99)", chi_squared);

    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 device = Device::Cpu;
    let q_values = Tensor::new(&[1.0f32, 2.0, 3.0], &device)?;

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

    // 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(())
}