//! Temperature Decay Unit Tests (Wave 2 Agent 2D)
//!
//! Tests for temperature decay rate calculation bug fix.
//!
//! Bug: Temperature reaches minimum too fast (459 epochs vs 1150 needed)
//! Target: Temperature should reach 0.1 minimum at ~75% of training (not 45%)
//! Fix: Calculate optimal decay rate based on target_temperature_epochs

#![allow(unused_crate_dependencies)]

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

#[test]
fn test_temperature_decay_reaches_min_at_75_percent() -> Result<()> {
    // Test that temperature reaches minimum at 75% of training
    let total_epochs = 1000;
    let target_pct = 0.75;
    let target_epoch = (total_epochs as f64 * target_pct) as usize;

    let mut config = WorkingDQNConfig::emergency_safe_defaults();
    config.temperature_start = 1.0;
    config.temperature_min = 0.1;

    // Calculate optimal decay rate for 75% convergence
    // decay = (temp_min / temp_start) ^ (1 / target_epochs)
    let target_epochs = (total_epochs as f64 * target_pct) as usize;
    config.temperature_decay =
        (config.temperature_min / config.temperature_start).powf(1.0 / target_epochs as f64);

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

    // Simulate temperature decay for target_epoch epochs
    for _ in 0..target_epoch {
        dqn.update_temperature();
    }

    // Temperature should be close to minimum at target epoch
    let temp_at_target = dqn.get_temperature();
    assert!(
        (temp_at_target - 0.1).abs() < 0.01,
        "Temperature should be ~0.1 at epoch {} (75% of training), got {:.6}",
        target_epoch,
        temp_at_target
    );

    // Continue to end of training
    for _ in target_epoch..total_epochs {
        dqn.update_temperature();
    }

    // Temperature should be at minimum (clamped)
    let final_temp = dqn.get_temperature();
    assert_eq!(
        final_temp, 0.1,
        "Temperature should be clamped at minimum 0.1 after {} epochs, got {:.6}",
        total_epochs, final_temp
    );

    Ok(())
}

#[test]
fn test_temperature_decay_calculation_different_epochs() -> Result<()> {
    // Test decay calculation with different epoch counts
    let test_cases = vec![
        (100, 0.75, 0.9698),  // 100 epochs, 75% convergence (75 target epochs)
        (500, 0.75, 0.9939),  // 500 epochs, 75% convergence (375 target epochs)
        (1000, 0.75, 0.9969), // 1000 epochs, 75% convergence (750 target epochs)
        (2000, 0.75, 0.9985), // 2000 epochs, 75% convergence (1500 target epochs)
    ];

    for (total_epochs, target_pct, expected_decay) in test_cases {
        let target_epochs = (total_epochs as f64 * target_pct) as usize;
        let temp_start = 1.0_f64;
        let temp_min = 0.1_f64;

        // Calculate decay: (temp_min / temp_start) ^ (1 / target_epochs)
        let calculated_decay = (temp_min / temp_start).powf(1.0 / target_epochs as f64);

        assert!(
            (calculated_decay - expected_decay).abs() < 0.0001,
            "Decay calculation for {} epochs should be ~{:.4}, got {:.6}",
            total_epochs,
            expected_decay,
            calculated_decay
        );
    }

    Ok(())
}

#[test]
fn test_temperature_decay_backward_compatibility() -> Result<()> {
    // Test that default config still works (backward compatibility)
    let mut config = WorkingDQNConfig::emergency_safe_defaults();
    config.temperature_start = 1.0;
    config.temperature_min = 0.1;
    config.temperature_decay = 0.995; // Old hardcoded value

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

    // Decay for 10 epochs
    for _ in 0..10 {
        dqn.update_temperature();
    }

    let temp_after_10 = dqn.get_temperature();
    let expected_temp = 1.0 * 0.995_f64.powi(10);

    assert!(
        (temp_after_10 - expected_temp).abs() < 0.0001,
        "Temperature after 10 epochs should be ~{:.6}, got {:.6}",
        expected_temp,
        temp_after_10
    );

    Ok(())
}

#[test]
fn test_temperature_never_goes_below_minimum() -> Result<()> {
    // Test that temperature is clamped at minimum
    let mut config = WorkingDQNConfig::emergency_safe_defaults();
    config.temperature_start = 1.0;
    config.temperature_min = 0.1;
    config.temperature_decay = 0.95; // Aggressive decay for testing

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

    // Decay for many epochs (well beyond convergence)
    for _ in 0..100 {
        dqn.update_temperature();
    }

    let final_temp = dqn.get_temperature();
    assert_eq!(
        final_temp, 0.1,
        "Temperature should be clamped at minimum 0.1, got {:.6}",
        final_temp
    );

    // Continue decaying - should stay at minimum
    for _ in 0..100 {
        dqn.update_temperature();
    }

    let still_min = dqn.get_temperature();
    assert_eq!(
        still_min, 0.1,
        "Temperature should remain at minimum 0.1, got {:.6}",
        still_min
    );

    Ok(())
}

#[test]
fn test_temperature_decay_rate_vs_epochs() -> Result<()> {
    // Verify the OLD bug: 0.995 decay reaches 0.1 after 459 epochs (45.9% of 1000)
    let mut config = WorkingDQNConfig::emergency_safe_defaults();
    config.temperature_start = 1.0;
    config.temperature_min = 0.1;
    config.temperature_decay = 0.995; // Old hardcoded value (buggy)

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

    let mut epoch_at_min = 0;
    for epoch in 0..1000 {
        dqn.update_temperature();
        if dqn.get_temperature() <= 0.1001 && epoch_at_min == 0 {
            epoch_at_min = epoch + 1; // +1 because epoch is 0-indexed
            break;
        }
    }

    // OLD BUG: Should reach minimum around epoch 459 (45.9% of 1000)
    assert!(
        epoch_at_min >= 450 && epoch_at_min <= 470,
        "Old decay (0.995) should reach minimum around epoch 459, got epoch {}",
        epoch_at_min
    );

    // Verify this is TOO EARLY (should be ~750 for 75% convergence)
    let expected_epoch_75pct = 750;
    let deviation = (epoch_at_min as i32 - expected_epoch_75pct as i32).abs();
    assert!(
        deviation > 200,
        "Old decay reaches minimum {} epochs too early (expected ~{}, got {})",
        deviation,
        expected_epoch_75pct,
        epoch_at_min
    );

    Ok(())
}

#[test]
fn test_optimal_decay_rate_formula() -> Result<()> {
    // Test the formula: decay = (temp_min / temp_start) ^ (1 / target_epochs)
    let total_epochs = 1000;
    let target_pct = 0.75;
    let target_epochs = (total_epochs as f64 * target_pct) as usize; // 750

    let temp_start = 1.0_f64;
    let temp_min = 0.1_f64;

    // Calculate optimal decay
    let optimal_decay = (temp_min / temp_start).powf(1.0 / target_epochs as f64);

    // Expected: (0.1 / 1.0) ^ (1 / 750) = 0.1 ^ (1/750) = 0.9969...
    assert!(
        (optimal_decay - 0.9969).abs() < 0.0001,
        "Optimal decay for 750 epochs should be ~0.9969, got {:.6}",
        optimal_decay
    );

    // Verify it reaches minimum at target epoch
    let mut config = WorkingDQNConfig::emergency_safe_defaults();
    config.temperature_start = temp_start;
    config.temperature_min = temp_min;
    config.temperature_decay = optimal_decay;

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

    for _ in 0..target_epochs {
        dqn.update_temperature();
    }

    let temp_at_target = dqn.get_temperature();
    assert!(
        (temp_at_target - temp_min).abs() < 0.01,
        "Temperature should be ~{} at epoch {}, got {:.6}",
        temp_min,
        target_epochs,
        temp_at_target
    );

    Ok(())
}