"""Mācīšanās ātruma plānotājs."""

from __future__ import annotations

import math


def cosine_schedule(step: int, total_steps: int, lr_max: float, lr_min: float = 1e-7) -> float:
    """Kosinusa LR grafiks."""
    progress = step / max(total_steps, 1)
    return lr_min + (lr_max - lr_min) * 0.5 * (1 + math.cos(math.pi * progress))


def warmup_schedule(step: int, warmup_steps: int, lr_max: float) -> float:
    """Lineārs sildīšanās grafiks."""
    if step >= warmup_steps:
        return lr_max
    return lr_max * step / max(warmup_steps, 1)