//! Sigma Annealing Scheduler for Noisy Networks
//!
//! Implements gradual annealing of noise standard deviation from initial
//! high exploration (`σ_init` = 0.6) to final refined exploration (`σ_final` = 0.4).
//!
//! This helps noisy networks transition from aggressive exploration early in training
//! to more refined exploitation later, improving convergence stability.

use serde::{Deserialize, Serialize};

/// Scheduler for annealing noisy network sigma parameters
///
/// # Example
/// ```
/// use foxhunt_ml::dqn::noisy_sigma_scheduler::NoisySigmaScheduler;
///
/// let mut scheduler = NoisySigmaScheduler::new(0.6, 0.4, 10000);
///
/// // At step 0
/// assert_eq!(scheduler.get_sigma(), 0.6);
///
/// // At step 5000 (50% progress)
/// scheduler.step_to(5000);
/// assert!((scheduler.get_sigma() - 0.5).abs() < 1e-6);
///
/// // At step 10000+ (100% progress)
/// scheduler.step_to(10000);
/// assert_eq!(scheduler.get_sigma(), 0.4);
/// ```
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NoisySigmaScheduler {
    /// Initial sigma value (e.g., 0.6 for aggressive early exploration)
    initial_sigma: f64,
    /// Final sigma value (e.g., 0.4 for refined late exploration)
    final_sigma: f64,
    /// Number of training steps over which to anneal
    decay_steps: usize,
    /// Current training step
    current_step: usize,
}

impl NoisySigmaScheduler {
    /// Create new sigma scheduler
    ///
    /// # Arguments
    /// * `initial_sigma` - Starting sigma value (typically 0.6)
    /// * `final_sigma` - Ending sigma value (typically 0.4)
    /// * `decay_steps` - Number of steps to anneal over
    ///
    /// # Example
    /// ```
    /// use foxhunt_ml::dqn::noisy_sigma_scheduler::NoisySigmaScheduler;
    ///
    /// let scheduler = NoisySigmaScheduler::new(0.6, 0.4, 10000);
    /// ```
    pub fn new(initial_sigma: f64, final_sigma: f64, decay_steps: usize) -> Self {
        assert!(
            initial_sigma >= final_sigma,
            "initial_sigma must be >= final_sigma (got {} < {})",
            initial_sigma,
            final_sigma
        );
        assert!(decay_steps > 0, "decay_steps must be > 0");

        Self {
            initial_sigma,
            final_sigma,
            decay_steps,
            current_step: 0,
        }
    }

    /// Get current sigma value based on linear annealing schedule
    ///
    /// Formula: σ(t) = `σ_init` × (1 - α) + `σ_final` × α
    /// where α = min(t / `decay_steps`, 1.0)
    pub fn get_sigma(&self) -> f64 {
        let progress = (self.current_step as f64 / self.decay_steps as f64).min(1.0);
        let annealed = self.initial_sigma * (1.0 - progress) + self.final_sigma * progress;
        annealed.max(Self::SIGMA_FLOOR)
    }

    /// Minimum sigma to prevent complete collapse of exploration noise
    const SIGMA_FLOOR: f64 = 0.3;

    /// Increment step counter
    pub const fn step(&mut self) {
        self.current_step += 1;
    }

    /// Set step counter to specific value (useful for resuming from checkpoint)
    pub const fn step_to(&mut self, step: usize) {
        self.current_step = step;
    }

    /// Reset scheduler to initial state
    pub const fn reset(&mut self) {
        self.current_step = 0;
    }

    /// Get current step
    pub const fn current_step(&self) -> usize {
        self.current_step
    }

    /// Check if annealing is complete
    pub const fn is_complete(&self) -> bool {
        self.current_step >= self.decay_steps
    }
}

impl Default for NoisySigmaScheduler {
    /// Default scheduler with wider initial exploration for 45-action space
    /// (0.8 → 0.4 over 100k steps, floor=0.3)
    fn default() -> Self {
        Self::new(0.8, 0.4, 100_000)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_scheduler_creation() {
        let scheduler = NoisySigmaScheduler::new(0.6, 0.4, 10000);
        assert_eq!(scheduler.initial_sigma, 0.6);
        assert_eq!(scheduler.final_sigma, 0.4);
        assert_eq!(scheduler.decay_steps, 10000);
        assert_eq!(scheduler.current_step, 0);
    }

    #[test]
    #[should_panic(expected = "initial_sigma must be >= final_sigma")]
    fn test_scheduler_invalid_sigma_range() {
        NoisySigmaScheduler::new(0.4, 0.6, 10000); // initial < final (invalid)
    }

    #[test]
    #[should_panic(expected = "decay_steps must be > 0")]
    fn test_scheduler_zero_decay_steps() {
        NoisySigmaScheduler::new(0.6, 0.4, 0);
    }

    #[test]
    fn test_sigma_at_start() {
        let scheduler = NoisySigmaScheduler::new(0.6, 0.4, 10000);
        assert_eq!(scheduler.get_sigma(), 0.6);
    }

    #[test]
    fn test_sigma_at_end() {
        let mut scheduler = NoisySigmaScheduler::new(0.6, 0.4, 10000);
        scheduler.step_to(10000);
        assert_eq!(scheduler.get_sigma(), 0.4);
    }

    #[test]
    fn test_sigma_at_midpoint() {
        let mut scheduler = NoisySigmaScheduler::new(0.6, 0.4, 10000);
        scheduler.step_to(5000); // 50% progress
        let expected = 0.5; // (0.6 + 0.4) / 2
        assert!((scheduler.get_sigma() - expected).abs() < 1e-6);
    }

    #[test]
    fn test_sigma_linear_interpolation() {
        let mut scheduler = NoisySigmaScheduler::new(0.6, 0.4, 10000);

        // 25% progress
        scheduler.step_to(2500);
        let expected = 0.6 * 0.75 + 0.4 * 0.25; // 0.55
        assert!((scheduler.get_sigma() - expected).abs() < 1e-6);

        // 75% progress
        scheduler.step_to(7500);
        let expected = 0.6 * 0.25 + 0.4 * 0.75; // 0.45
        assert!((scheduler.get_sigma() - expected).abs() < 1e-6);
    }

    #[test]
    fn test_sigma_beyond_decay_steps() {
        let mut scheduler = NoisySigmaScheduler::new(0.6, 0.4, 10000);
        scheduler.step_to(15000); // Beyond decay_steps
        assert_eq!(scheduler.get_sigma(), 0.4); // Clamped to final
    }

    #[test]
    fn test_step_increments() {
        let mut scheduler = NoisySigmaScheduler::new(0.6, 0.4, 100);

        for i in 0..100 {
            assert_eq!(scheduler.current_step(), i);
            scheduler.step();
        }

        assert_eq!(scheduler.current_step(), 100);
        assert!(scheduler.is_complete());
    }

    #[test]
    fn test_step_to() {
        let mut scheduler = NoisySigmaScheduler::new(0.6, 0.4, 10000);

        scheduler.step_to(3000);
        assert_eq!(scheduler.current_step(), 3000);

        scheduler.step_to(7000);
        assert_eq!(scheduler.current_step(), 7000);
    }

    #[test]
    fn test_reset() {
        let mut scheduler = NoisySigmaScheduler::new(0.6, 0.4, 10000);

        scheduler.step_to(5000);
        assert_eq!(scheduler.current_step(), 5000);

        scheduler.reset();
        assert_eq!(scheduler.current_step(), 0);
        assert_eq!(scheduler.get_sigma(), 0.6);
    }

    #[test]
    fn test_is_complete() {
        let mut scheduler = NoisySigmaScheduler::new(0.6, 0.4, 100);

        assert!(!scheduler.is_complete());

        scheduler.step_to(50);
        assert!(!scheduler.is_complete());

        scheduler.step_to(100);
        assert!(scheduler.is_complete());

        scheduler.step_to(150);
        assert!(scheduler.is_complete());
    }

    #[test]
    fn test_default_scheduler() {
        let scheduler = NoisySigmaScheduler::default();
        assert_eq!(scheduler.initial_sigma, 0.8);
        assert_eq!(scheduler.final_sigma, 0.4);
        assert_eq!(scheduler.decay_steps, 100_000);
    }

    #[test]
    fn test_sigma_floor() {
        // Even if final_sigma is below floor, get_sigma() respects the floor
        let mut scheduler = NoisySigmaScheduler::new(0.5, 0.1, 100);
        scheduler.step_to(100);
        assert!(scheduler.get_sigma() >= NoisySigmaScheduler::SIGMA_FLOOR,
            "Sigma {} should be >= floor {}", scheduler.get_sigma(), NoisySigmaScheduler::SIGMA_FLOOR);
    }

    #[test]
    fn test_monotonic_decrease() {
        let mut scheduler = NoisySigmaScheduler::new(0.6, 0.4, 1000);

        let mut prev_sigma = scheduler.get_sigma();
        for _ in 0..1000 {
            scheduler.step();
            let current_sigma = scheduler.get_sigma();
            assert!(
                current_sigma <= prev_sigma,
                "Sigma should monotonically decrease: {} > {}",
                current_sigma,
                prev_sigma
            );
            prev_sigma = current_sigma;
        }
    }

    #[test]
    fn test_serialization() {
        let scheduler = NoisySigmaScheduler::new(0.6, 0.4, 50000);

        let json = serde_json::to_string(&scheduler).unwrap();
        let deserialized: NoisySigmaScheduler = serde_json::from_str(&json).unwrap();

        assert_eq!(deserialized.initial_sigma, scheduler.initial_sigma);
        assert_eq!(deserialized.final_sigma, scheduler.final_sigma);
        assert_eq!(deserialized.decay_steps, scheduler.decay_steps);
        assert_eq!(deserialized.current_step, scheduler.current_step);
    }
}
