//! Bug #28: Unused import warning for Device in softmax.rs
//!
//! This test ensures ml/src/dqn/softmax.rs compiles without warnings
//! after fixing the conditional compilation of Device import.

use ml::dqn::softmax::{softmax_with_temperature, sample_from_softmax, softmax_entropy};
use candle_core::{Device, Tensor};

#[test]
fn test_bug28_softmax_imports_clean() {
    // Verify softmax.rs compiles without unused import warnings
    // This test exercises all public functions to ensure they work correctly

    let device = Device::Cpu;

    // Test 1: softmax_with_temperature
    let q_values = Tensor::new(&[1.0f32, 2.0, 3.0], &device)
        .expect("Failed to create Q-values tensor");

    let probs = softmax_with_temperature(&q_values, 1.0)
        .expect("Failed to compute softmax");

    let probs_vec = probs.to_vec1::<f32>()
        .expect("Failed to convert probabilities to vector");

    // Verify probabilities 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
    );

    // Test 2: sample_from_softmax
    let action = sample_from_softmax(&q_values, 1.0)
        .expect("Failed to sample action");

    assert!(action < 3, "Action should be in range [0, 3), got {}", action);

    // Test 3: softmax_entropy
    let entropy = softmax_entropy(&q_values, 1.0)
        .expect("Failed to compute entropy");

    assert!(entropy > 0.0, "Entropy should be positive, got {}", entropy);
    assert!(entropy < 2.0, "Entropy should be < log2(3) ≈ 1.585, got {}", entropy);
}

#[test]
fn test_bug28_softmax_numerical_stability() {
    // Test the log-sum-exp trick for numerical stability
    let device = Device::Cpu;

    // Large Q-values that would overflow without stability trick
    let q_values = Tensor::new(&[100.0f32, 200.0, 300.0], &device)
        .expect("Failed to create large Q-values");

    let probs = softmax_with_temperature(&q_values, 1.0)
        .expect("Failed to compute softmax with large values");

    let probs_vec = probs.to_vec1::<f32>()
        .expect("Failed to convert probabilities to vector");

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

    // Verify no NaN or Inf
    assert!(probs_vec.iter().all(|&p| p.is_finite()), "All probabilities should be finite");
}

#[test]
fn test_bug28_temperature_control() {
    // Test temperature parameter effect on exploration
    let device = Device::Cpu;
    let q_values = Tensor::new(&[1.0f32, 2.0, 3.0], &device)
        .expect("Failed to create Q-values");

    // Low temperature (greedy)
    let entropy_low = softmax_entropy(&q_values, 0.1)
        .expect("Failed to compute entropy with low temp");

    // High temperature (exploratory)
    let entropy_high = softmax_entropy(&q_values, 10.0)
        .expect("Failed to compute entropy with high temp");

    // High temperature should produce higher entropy
    assert!(
        entropy_high > entropy_low,
        "High temperature ({}) should produce higher entropy than low temperature ({})",
        entropy_high,
        entropy_low
    );
}