// Author: Simon-Pierre Boucher — contact@spboucher.ai
#pragma once

#include <cmath>
#include <cstdint>

namespace forge::train {

// Linear warmup → cosine decay to min_lr (nanoGPT schedule; (it+1)/(warmup+1)
// avoids a zero-lr first step).
inline float lr_at(int64_t step, float max_lr, float min_lr, int64_t warmup_steps,
                   int64_t decay_steps) {
    if (step < warmup_steps) return max_lr * float(step + 1) / float(warmup_steps + 1);
    if (step >= decay_steps) return min_lr;
    const float ratio = float(step - warmup_steps) / float(decay_steps - warmup_steps);
    const float coeff = 0.5f * (1.0f + std::cos(float(M_PI) * ratio));
    return min_lr + coeff * (max_lr - min_lr);
}

// Warmup–Stable–Decay (MiniCPM): linear warmup, flat plateau at max_lr, then a
// short cooldown over the final decay_steps. Any plateau checkpoint can be
// resumed and the run extended without re-deciding total_steps up front. The
// cooldown uses the 1-sqrt shape, which beats linear/exponential cooldowns in
// the WSD ablations (arXiv 2404.06395 follow-ups).
inline float lr_wsd(int64_t step, float max_lr, float min_lr, int64_t warmup_steps,
                    int64_t total_steps, int64_t decay_steps) {
    if (step < warmup_steps) return max_lr * float(step + 1) / float(warmup_steps + 1);
    const int64_t decay_start = total_steps - decay_steps;
    if (step < decay_start) return max_lr;
    if (step >= total_steps) return min_lr;
    const float ratio = float(step - decay_start) / float(decay_steps);
    return min_lr + (max_lr - min_lr) * (1.0f - std::sqrt(ratio));
}

} // namespace forge::train