"""Learning-rate schedule: linear warmup, then cosine decay to a floor."""

from __future__ import annotations

import math


def lr_at(step: int, peak_lr: float, warmup_steps: int, total_steps: int, min_ratio: float = 0.1) -> float:
    if step < warmup_steps:
        return peak_lr * (step + 1) / warmup_steps
    progress = min(1.0, (step - warmup_steps) / max(1, total_steps - 1 - warmup_steps))
    floor = peak_lr * min_ratio
    return floor + (peak_lr - floor) * 0.5 * (1 + math.cos(math.pi * progress))