from __future__ import annotations

from dataclasses import dataclass

from sandevistanlr.core.schedule_base import SandevistanScheduleBase


@dataclass
class SandevistanLR:
    """
    PyTorch-compatible learning-rate callable.

    Example:
        base = SandevistanScheduleBase(...)
        lr_fn = SandevistanLR(base)
        lr = lr_fn(global_step)
    """

    base: SandevistanScheduleBase

    def __call__(self, step: int) -> float:
        return float(self.base.lr_at_step(step))
