//! Wave 3 Agent 2: Temperature Floor Adjustment Test
//!
//! Verifies that temperature_min=0.3 prevents over-exploitation in late training.
//!
//! Key Requirements:
//! 1. Temperature reaches 0.3 (not 0.1) at target epoch fraction
//! 2. Temperature never goes below 0.3
//! 3. Temperature provides ~70% probability on max Q-value (vs 99% at 0.1)

use ml::dqn::{WorkingDQN, WorkingDQNConfig};

#[test]
fn test_temperature_floor_0_3() -> anyhow::Result<()> {
    // Setup: 50-epoch training scenario
    let mut config = WorkingDQNConfig::emergency_safe_defaults();
    config.temperature_start = 1.0;
    config.temperature_min = 0.3; // NEW FLOOR
    config.target_temperature_fraction = 0.75; // 75% of training

    // Calculate optimal decay to reach 0.3 at epoch 37 (75% of 50 epochs)
    let total_epochs = 50;
    config.temperature_decay = WorkingDQNConfig::calculate_optimal_temperature_decay(
        total_epochs,
        config.temperature_start,
        config.temperature_min,
        config.target_temperature_fraction,
    );

    let mut dqn = WorkingDQN::new(config)?;

    // Verify initial temperature
    assert_eq!(
        dqn.get_temperature(),
        1.0,
        "Initial temperature should be 1.0"
    );

    // Simulate 37 epochs of temperature decay (75% of 50 epochs)
    for epoch in 1..=37 {
        dqn.update_temperature();

        // Temperature should never go below 0.3
        assert!(
            dqn.get_temperature() >= 0.3,
            "Temperature at epoch {} is {} (below floor 0.3)",
            epoch,
            dqn.get_temperature()
        );
    }

    // At epoch 37 (75% of 50-epoch training), temperature should be at or near floor
    let temp_at_target = dqn.get_temperature();
    assert!(
        temp_at_target >= 0.3,
        "Temperature at target epoch (37) is {} (below floor 0.3)",
        temp_at_target
    );

    // Continue decaying for another 50 epochs - should stay at 0.3
    for epoch in 38..=87 {
        dqn.update_temperature();
        assert_eq!(
            dqn.get_temperature(),
            0.3,
            "Temperature at epoch {} is {} (should be clamped at 0.3)",
            epoch,
            dqn.get_temperature()
        );
    }

    Ok(())
}

#[test]
fn test_temperature_never_below_floor() -> anyhow::Result<()> {
    // Edge case: very aggressive decay should still respect floor
    let mut config = WorkingDQNConfig::emergency_safe_defaults();
    config.temperature_start = 1.0;
    config.temperature_min = 0.3;
    config.temperature_decay = 0.9; // Aggressive decay

    let mut dqn = WorkingDQN::new(config)?;

    // Decay for 100 epochs
    for _ in 0..100 {
        dqn.update_temperature();
        assert!(
            dqn.get_temperature() >= 0.3,
            "Temperature {} below floor 0.3 with aggressive decay",
            dqn.get_temperature()
        );
    }

    // Should be exactly 0.3 after clamping
    assert_eq!(dqn.get_temperature(), 0.3);

    Ok(())
}

#[test]
fn test_optimal_decay_calculation() -> anyhow::Result<()> {
    // Test that calculate_optimal_temperature_decay produces correct convergence
    let total_epochs = 50;
    let temp_start = 1.0;
    let temp_min = 0.3;
    let target_fraction = 0.75; // 75% = 37.5 epochs

    let optimal_decay = WorkingDQNConfig::calculate_optimal_temperature_decay(
        total_epochs,
        temp_start,
        temp_min,
        target_fraction,
    );

    // Verify decay reaches minimum at target epoch
    let target_epochs = (total_epochs as f64 * target_fraction) as usize;
    let mut temperature = temp_start;

    for _ in 0..target_epochs {
        temperature *= optimal_decay;
        temperature = temperature.max(temp_min);
    }

    // Temperature should be at or very close to minimum (within 1% tolerance)
    let tolerance = temp_min * 0.01;
    assert!(
        (temperature - temp_min).abs() < tolerance,
        "Temperature {} not close to minimum {} after {} epochs (tolerance {})",
        temperature,
        temp_min,
        target_epochs,
        tolerance
    );

    Ok(())
}

#[test]
fn test_temperature_floor_prevents_overexploitation() -> anyhow::Result<()> {
    // Verify that temp=0.3 provides better exploration than temp=0.1
    // Softmax: p(a) = exp(Q(a)/T) / sum(exp(Q(i)/T))
    // With Q-values: [1.0, 0.5, 0.3] (BUY best)

    let q_values = vec![1.0, 0.5, 0.3];

    // Calculate softmax probabilities at temp=0.3 (new floor)
    let temp_03: f64 = 0.3;
    let logits_03: Vec<f64> = q_values.iter().map(|q| (*q / temp_03).exp()).collect();
    let sum_03: f64 = logits_03.iter().sum();
    let probs_03: Vec<f64> = logits_03.iter().map(|l| l / sum_03).collect();

    // Calculate softmax probabilities at temp=0.1 (old floor)
    let temp_01: f64 = 0.1;
    let logits_01: Vec<f64> = q_values.iter().map(|q| (*q / temp_01).exp()).collect();
    let sum_01: f64 = logits_01.iter().sum();
    let probs_01: Vec<f64> = logits_01.iter().map(|l| l / sum_01).collect();

    // At temp=0.3, max action should have ~70% probability (balanced)
    assert!(
        probs_03[0] >= 0.60 && probs_03[0] <= 0.80,
        "Temp=0.3 probability for max Q-value is {:.2}% (expected ~70%)",
        probs_03[0] * 100.0
    );

    // At temp=0.1, max action should have ~99% probability (over-exploitation)
    assert!(
        probs_01[0] >= 0.95,
        "Temp=0.1 probability for max Q-value is {:.2}% (expected >95%)",
        probs_01[0] * 100.0
    );

    // Verify temp=0.3 provides more exploration (>20% on non-max actions)
    let exploration_03 = 1.0 - probs_03[0];
    let exploration_01 = 1.0 - probs_01[0];

    assert!(
        exploration_03 > exploration_01 * 5.0,
        "Temp=0.3 exploration ({:.2}%) not significantly higher than temp=0.1 ({:.2}%)",
        exploration_03 * 100.0,
        exploration_01 * 100.0
    );

    Ok(())
}

#[test]
fn test_default_config_uses_new_floor() -> anyhow::Result<()> {
    // Verify that default config now uses 0.3 instead of 0.1
    let config = WorkingDQNConfig::emergency_safe_defaults();

    assert_eq!(
        config.temperature_min, 0.3,
        "Default config should use temperature_min=0.3, got {}",
        config.temperature_min
    );

    Ok(())
}