//! Learning Rate Scheduler with Warmup
//!
//! Implements linear warmup followed by configurable decay strategies.
//! Pattern copied from TFT trainer (ml/src/tft/training.rs).
//!
//! ## Features
//! - Linear warmup from 0 to initial_lr over warmup_steps
//! - Constant LR after warmup
//! - Linear decay from initial_lr to end_lr
//! - Cosine annealing decay from initial_lr to min_lr
//! - Exponential decay from initial_lr to min_lr
//!
//! ## Usage
//! ```rust,ignore
//! let scheduler = LRScheduler::new(
//!     0.001,  // initial_lr
//!     1000,   // warmup_steps
//!     LRDecayType::Cosine { min_lr: 1e-6, total_steps: 10000 }
//! );
//!
//! for epoch in 0..10000 {
//!     let lr = scheduler.get_lr();
//!     optimizer.set_learning_rate(lr);
//!     // ... training step ...
//!     scheduler.step();
//! }
//! ```

use serde::{Deserialize, Serialize};

/// Learning rate decay type after warmup
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq)]
pub enum LRDecayType {
    /// Constant learning rate (no decay)
    Constant,
    /// Linear decay from initial_lr to end_lr over total_steps
    Linear {
        end_lr: f64,
        total_steps: usize,
    },
    /// Cosine annealing from initial_lr to min_lr over total_steps
    Cosine {
        min_lr: f64,
        total_steps: usize,
    },
    /// Exponential decay: lr = max(initial_lr * decay_rate^step, min_lr)
    Exponential {
        decay_rate: f64,
        min_lr: f64,
    },
}

/// Learning rate scheduler with linear warmup and configurable decay
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LRScheduler {
    /// Initial learning rate (target after warmup)
    initial_lr: f64,
    /// Number of warmup steps (linear ramp from 0 to initial_lr)
    warmup_steps: usize,
    /// Decay type after warmup
    decay_type: LRDecayType,
    /// Current step counter
    current_step: usize,
}

impl LRScheduler {
    /// Create new learning rate scheduler
    ///
    /// # Arguments
    /// * `initial_lr` - Target learning rate after warmup
    /// * `warmup_steps` - Number of steps for linear warmup (0 = no warmup)
    /// * `decay_type` - Decay strategy after warmup
    pub fn new(initial_lr: f64, warmup_steps: usize, decay_type: LRDecayType) -> Self {
        Self {
            initial_lr,
            warmup_steps,
            decay_type,
            current_step: 0,
        }
    }

    /// Get current learning rate
    pub fn get_lr(&self) -> f64 {
        // Phase 1: Linear warmup
        if self.current_step < self.warmup_steps {
            if self.warmup_steps == 0 {
                return self.initial_lr;
            }
            // Linear interpolation from 0 to initial_lr
            return self.initial_lr * (self.current_step as f64 / self.warmup_steps as f64);
        }

        // Phase 2: Apply decay
        let steps_after_warmup = self.current_step - self.warmup_steps;

        match self.decay_type {
            LRDecayType::Constant => self.initial_lr,

            LRDecayType::Linear { end_lr, total_steps } => {
                // Calculate decay steps (total_steps includes warmup)
                let decay_steps = total_steps.saturating_sub(self.warmup_steps);

                if decay_steps == 0 || steps_after_warmup >= decay_steps {
                    return end_lr;
                }

                // Linear interpolation from initial_lr to end_lr
                let progress = steps_after_warmup as f64 / decay_steps as f64;
                self.initial_lr - (self.initial_lr - end_lr) * progress
            }

            LRDecayType::Cosine { min_lr, total_steps } => {
                // Calculate decay steps (total_steps includes warmup)
                let decay_steps = total_steps.saturating_sub(self.warmup_steps);

                if decay_steps == 0 || steps_after_warmup >= decay_steps {
                    return min_lr;
                }

                // Cosine annealing from initial_lr to min_lr
                let progress = steps_after_warmup as f64 / decay_steps as f64;
                min_lr + (self.initial_lr - min_lr)
                    * (1.0 + (std::f64::consts::PI * progress).cos()) / 2.0
            }

            LRDecayType::Exponential { decay_rate, min_lr } => {
                // Exponential decay: lr = max(initial_lr * decay_rate^step, min_lr)
                let decayed_lr = self.initial_lr * decay_rate.powf(steps_after_warmup as f64);
                decayed_lr.max(min_lr)
            }
        }
    }

    /// Advance to next step (call after each training step)
    pub fn step(&mut self) {
        self.current_step += 1;
    }

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

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

    /// Get initial learning rate
    pub fn get_initial_lr(&self) -> f64 {
        self.initial_lr
    }

    /// Get warmup steps
    pub fn get_warmup_steps(&self) -> usize {
        self.warmup_steps
    }

    /// Get decay type
    pub fn get_decay_type(&self) -> LRDecayType {
        self.decay_type
    }
}

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

    #[test]
    fn test_constant_lr() {
        let scheduler = LRScheduler::new(0.001, 0, LRDecayType::Constant);
        assert_eq!(scheduler.get_lr(), 0.001);
    }

    #[test]
    fn test_warmup() {
        let scheduler = LRScheduler::new(0.001, 10, LRDecayType::Constant);
        assert_eq!(scheduler.get_lr(), 0.0); // Step 0: 0%
    }

    #[test]
    fn test_linear_decay() {
        let scheduler = LRScheduler::new(
            0.001,
            0,
            LRDecayType::Linear { end_lr: 0.0001, total_steps: 100 }
        );
        assert_eq!(scheduler.get_lr(), 0.001);
    }

    #[test]
    fn test_cosine_decay() {
        let scheduler = LRScheduler::new(
            0.001,
            0,
            LRDecayType::Cosine { min_lr: 0.0001, total_steps: 100 }
        );
        assert_eq!(scheduler.get_lr(), 0.001);
    }

    #[test]
    fn test_exponential_decay() {
        let mut scheduler = LRScheduler::new(
            0.001,
            0,
            LRDecayType::Exponential { decay_rate: 0.95, min_lr: 1e-6 }
        );

        // Step 0: initial_lr
        assert_eq!(scheduler.get_lr(), 0.001);

        // Step 1: initial_lr * decay_rate^1
        scheduler.step();
        let expected_lr_1 = 0.001 * 0.95;
        assert!((scheduler.get_lr() - expected_lr_1).abs() < 1e-9);

        // Step 2: initial_lr * decay_rate^2
        scheduler.step();
        let expected_lr_2 = 0.001 * 0.95_f64.powf(2.0);
        assert!((scheduler.get_lr() - expected_lr_2).abs() < 1e-9);

        // After many steps, should approach min_lr
        for _ in 0..1000 {
            scheduler.step();
        }
        assert!(scheduler.get_lr() >= 1e-6);
    }
}
