//! Learning-rate schedule: linear warmup → cosine decay to a floor. Host-only
//! and pure (just arithmetic over the step index), so it unit-tests locally.

/// Warmup-then-cosine LR schedule.
///
/// - steps `0..warmup`: linear ramp `0 → max_lr`
/// - steps `warmup..total`: cosine decay `max_lr → min_lr`
/// - steps `>= total`: clamped at `min_lr`
#[derive(Clone, Copy, Debug)]
pub struct LrSchedule {
    pub max_lr: f32,
    pub min_lr: f32,
    pub warmup: usize,
    pub total: usize,
}

impl LrSchedule {
    /// LR for a 0-indexed step.
    pub fn lr(&self, step: usize) -> f32 {
        if step < self.warmup {
            // Linear warmup; +1 so step 0 already takes a (tiny) nonzero LR.
            return self.max_lr * (step + 1) as f32 / self.warmup.max(1) as f32;
        }
        if step >= self.total {
            return self.min_lr;
        }
        let progress = (step - self.warmup) as f32 / (self.total - self.warmup).max(1) as f32;
        let cosine = 0.5 * (1.0 + (std::f32::consts::PI * progress).cos()); // 1 → 0
        self.min_lr + (self.max_lr - self.min_lr) * cosine
    }
}

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

    #[test]
    fn warmup_then_cosine_shape() {
        let s = LrSchedule {
            max_lr: 1.0,
            min_lr: 0.1,
            warmup: 10,
            total: 100,
        };
        // Warmup ramps up and reaches the peak at the end of warmup.
        assert!(s.lr(0) < s.lr(5));
        assert!(s.lr(5) < s.lr(9));
        assert!((s.lr(9) - 1.0).abs() < 1e-6);
        // Just after warmup, near the peak; midway, near the midpoint.
        assert!(s.lr(10) > 0.95);
        let mid = s.lr(55); // progress ~0.5 → cosine ~0.5
        assert!((mid - (0.1 + 0.9 * 0.5)).abs() < 0.02);
        // Decays monotonically to the floor by the end.
        assert!(s.lr(99) < s.lr(55));
        assert!(s.lr(100) <= s.min_lr + 1e-6);
        assert!(s.lr(1_000) <= s.min_lr + 1e-6);
    }
}
