//! Cosine annealing learning rate scheduler
//!
//! Implements cosine annealing with warm restarts for efficient transformer training.

use crate::schedulers::LearningRateScheduler;
use crate::{Result, TransformerError};
use serde::{Deserialize, Serialize};
use tracing::{debug, trace};

/// Cosine annealing learning rate scheduler
///
/// Decreases the learning rate following a cosine function, with optional warm restarts.
/// This provides a smooth decay that often leads to better convergence than step-wise decay.
///
/// # Mathematical Foundation
///
/// For step t in the current cycle:
/// - lr(t) = `min_lr` + (`base_lr` - `min_lr`) * 0.5 * (1 + cos(π * t / `T_max`))
///
/// Where `T_max` is the maximum number of steps in the current cycle.
///
/// # Benefits
/// - Smooth learning rate decay
/// - Can escape local minima with warm restarts
/// - Widely used in state-of-the-art transformer training
/// - Better final convergence than linear decay
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CosineAnnealingScheduler {
    /// Maximum learning rate at the beginning of each cycle
    base_lr: f64,
    /// Minimum learning rate at the end of each cycle
    min_lr: f64,
    /// Number of steps in each cycle
    t_max: usize,
    /// Current step count
    current_step: usize,
    /// Whether to use warm restarts
    use_restarts: bool,
}

impl CosineAnnealingScheduler {
    /// Create a new cosine annealing scheduler
    ///
    /// # Arguments
    /// * `base_lr` - Maximum learning rate (must be positive)
    /// * `min_lr` - Minimum learning rate (must be non-negative and <= `base_lr`)
    /// * `t_max` - Number of steps in each cycle (must be positive)
    ///
    /// # Errors
    /// Returns error if parameters are invalid
    pub fn new(base_lr: f64, min_lr: f64, t_max: usize) -> Result<Self> {
        Self::with_restarts(base_lr, min_lr, t_max, true)
    }

    /// Create a new cosine annealing scheduler with optional restarts
    ///
    /// # Arguments
    /// * `base_lr` - Maximum learning rate (must be positive)
    /// * `min_lr` - Minimum learning rate (must be non-negative and <= `base_lr`)
    /// * `t_max` - Number of steps in each cycle (must be positive)
    /// * `use_restarts` - Whether to restart the cycle after `T_max` steps
    ///
    /// # Errors
    /// Returns error if parameters are invalid
    pub fn with_restarts(
        base_lr: f64,
        min_lr: f64,
        t_max: usize,
        use_restarts: bool,
    ) -> Result<Self> {
        if base_lr <= 0.0 {
            return Err(TransformerError::generic(format!(
                "base_lr {base_lr} must be positive"
            )));
        }

        if min_lr < 0.0 {
            return Err(TransformerError::generic(format!(
                "min_lr {min_lr} must be non-negative"
            )));
        }

        if min_lr > base_lr {
            return Err(TransformerError::generic(format!(
                "min_lr {min_lr} must be <= base_lr {base_lr}"
            )));
        }

        if t_max == 0 {
            return Err(TransformerError::generic(format!(
                "t_max {t_max} must be positive"
            )));
        }

        debug!(
            "Creating cosine annealing scheduler: base_lr={}, min_lr={}, t_max={}, restarts={}",
            base_lr, min_lr, t_max, use_restarts
        );

        Ok(Self {
            base_lr,
            min_lr,
            t_max,
            current_step: 0,
            use_restarts,
        })
    }

    /// Get the minimum learning rate
    #[must_use]
    pub fn min_lr(&self) -> f64 {
        self.min_lr
    }

    /// Get the cycle length
    #[must_use]
    pub fn t_max(&self) -> usize {
        self.t_max
    }

    /// Check if using warm restarts
    #[must_use]
    pub fn use_restarts(&self) -> bool {
        self.use_restarts
    }

    /// Get the current cycle number (0-indexed)
    ///
    /// Each cycle contains `t_max` + 1 steps (from 0 to `t_max` inclusive).
    #[must_use]
    pub fn current_cycle(&self, step: usize) -> usize {
        if self.use_restarts {
            step / (self.t_max + 1)
        } else {
            0
        }
    }

    /// Get the step within the current cycle
    ///
    /// For a cycle of `t_max` steps:
    /// - Steps 0 to `t_max` form one complete cycle (`t_max` + 1 total steps)
    /// - At step `t_max`, LR reaches `min_lr`
    /// - At step `t_max` + 1, a new cycle begins at `base_lr`
    #[must_use]
    pub fn cycle_step(&self, step: usize) -> usize {
        if self.use_restarts {
            // Use t_max + 1 as period so step t_max is end of cycle (min_lr)
            // and step t_max + 1 starts a new cycle (base_lr)
            step % (self.t_max + 1)
        } else {
            step.min(self.t_max)
        }
    }

    /// Get the progress within the current cycle (0.0 to 1.0)
    #[must_use]
    pub fn cycle_progress(&self, step: usize) -> f64 {
        let cycle_step = self.cycle_step(step);
        cycle_step as f64 / self.t_max as f64
    }
}

impl LearningRateScheduler for CosineAnnealingScheduler {
    fn get_lr(&self, _epoch: usize, step: usize) -> f64 {
        let cycle_step = self.cycle_step(step);
        let progress = cycle_step as f64 / self.t_max as f64;

        // Cosine annealing formula
        let cosine_factor = 0.5 * (1.0 + (std::f64::consts::PI * progress).cos());
        let lr = self.min_lr + (self.base_lr - self.min_lr) * cosine_factor;

        trace!(
            "Cosine step {}: cycle_step={}, progress={:.4}, lr={:.6}",
            step, cycle_step, progress, lr
        );

        lr
    }

    fn step(&mut self) {
        self.current_step += 1;
        trace!(
            "Cosine annealing scheduler stepped to: {}",
            self.current_step
        );
    }

    fn current_step(&self) -> usize {
        self.current_step
    }

    fn reset(&mut self) {
        self.current_step = 0;
        debug!("Reset cosine annealing scheduler");
    }

    fn scheduler_type(&self) -> &'static str {
        if self.use_restarts {
            "CosineAnnealingWarmRestarts"
        } else {
            "CosineAnnealing"
        }
    }

    fn base_lr(&self) -> f64 {
        self.base_lr
    }
}

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

    #[test]
    fn test_cosine_annealing_creation() {
        let scheduler = CosineAnnealingScheduler::new(0.001, 0.00001, 1000).unwrap();
        assert_eq!(scheduler.base_lr(), 0.001);
        assert_eq!(scheduler.min_lr(), 0.00001);
        assert_eq!(scheduler.t_max(), 1000);
        assert_eq!(scheduler.current_step(), 0);
        assert!(scheduler.use_restarts());
    }

    #[test]
    fn test_cosine_annealing_invalid_params() {
        assert!(CosineAnnealingScheduler::new(-0.001, 0.00001, 1000).is_err());
        assert!(CosineAnnealingScheduler::new(0.001, -0.00001, 1000).is_err());
        assert!(CosineAnnealingScheduler::new(0.001, 0.002, 1000).is_err());
        assert!(CosineAnnealingScheduler::new(0.001, 0.00001, 0).is_err());
    }

    #[test]
    fn test_cosine_annealing_progression() {
        let scheduler = CosineAnnealingScheduler::new(0.001, 0.00001, 1000).unwrap();

        // At step 0, should be at base_lr
        let lr_0 = scheduler.get_lr(0, 0);
        assert!((lr_0 - 0.001).abs() < 1e-10);

        // At step t_max, should be at min_lr
        let lr_1000 = scheduler.get_lr(0, 1000);
        assert!((lr_1000 - 0.00001).abs() < 1e-6);
    }

    #[test]
    fn test_cosine_annealing_restarts() {
        let scheduler = CosineAnnealingScheduler::new(0.001, 0.00001, 100).unwrap();

        // At end of first cycle (step t_max = 100)
        let lr_100 = scheduler.get_lr(0, 100);
        assert!((lr_100 - 0.00001).abs() < 1e-6);

        // At start of second cycle (step t_max + 1 = 101)
        let lr_101 = scheduler.get_lr(0, 101);
        assert!((lr_101 - 0.001).abs() < 1e-6);

        // Cycle step verification:
        // - step 100 is at end of cycle 1 (cycle_step = 100)
        // - step 101 is at start of cycle 2 (cycle_step = 0)
        assert_eq!(scheduler.cycle_step(100), 100);
        assert_eq!(scheduler.cycle_step(101), 0);
        assert_eq!(scheduler.cycle_step(102), 1);
    }

    #[test]
    fn test_cosine_annealing_without_restarts() {
        let scheduler =
            CosineAnnealingScheduler::with_restarts(0.001, 0.00001, 100, false).unwrap();
        assert!(!scheduler.use_restarts());

        // After t_max, should stay at min_lr
        let lr_100 = scheduler.get_lr(0, 100);
        let lr_150 = scheduler.get_lr(0, 150);

        assert!((lr_100 - 0.00001).abs() < 1e-6);
        assert!((lr_150 - 0.00001).abs() < 1e-6);
    }

    #[test]
    fn test_cycle_calculations() {
        let scheduler = CosineAnnealingScheduler::new(0.001, 0.00001, 100).unwrap();

        // Each cycle has t_max + 1 = 101 steps (0 to 100 inclusive)
        // Cycle 0: steps 0-100
        // Cycle 1: steps 101-201
        // Cycle 2: steps 202-302
        assert_eq!(scheduler.current_cycle(0), 0);
        assert_eq!(scheduler.current_cycle(100), 0); // Still in cycle 0 (last step)
        assert_eq!(scheduler.current_cycle(101), 1); // Start of cycle 1
        assert_eq!(scheduler.current_cycle(201), 1); // Last step of cycle 1
        assert_eq!(scheduler.current_cycle(202), 2); // Start of cycle 2

        assert_eq!(scheduler.cycle_step(0), 0);
        assert_eq!(scheduler.cycle_step(99), 99);
        assert_eq!(scheduler.cycle_step(100), 100); // End of cycle 0
        assert_eq!(scheduler.cycle_step(101), 0); // Start of cycle 1
        assert_eq!(scheduler.cycle_step(150), 49); // 150 % 101 = 49

        assert_eq!(scheduler.cycle_progress(0), 0.0);
        assert_eq!(scheduler.cycle_progress(50), 0.5);
        assert_eq!(scheduler.cycle_progress(100), 1.0); // End of cycle
    }
}