from __future__ import annotations

from dataclasses import dataclass
from typing import Any

from sandevistanlr.core.schedule import SandevistanSchedule


@dataclass
class SandevistanKerasCallback:
    schedule: SandevistanSchedule
    verbose: int = 0

    def on_train_begin(self, logs: Any = None) -> None:
        self._step = 0

    def _set_lr(self, optimizer: Any, lr: float) -> None:
        if hasattr(optimizer, "learning_rate"):
            try:
                optimizer.learning_rate.assign(lr)
                return
            except Exception:
                optimizer.learning_rate = lr
                return

        if hasattr(optimizer, "lr"):
            try:
                optimizer.lr.assign(lr)
                return
            except Exception:
                optimizer.lr = lr
                return

        raise AttributeError("Optimizer has no learning_rate/lr attribute.")

    def on_train_batch_begin(self, batch: int, logs: Any = None) -> None:
        lr = float(self.schedule.value(self._step))
        self._set_lr(self.model.optimizer, lr)  # type: ignore[attr-defined]
        if self.verbose and (self._step % 100 == 0):
            print(f"[SandevistanLR] step={self._step} lr={lr:.6e}")
        self._step += 1


# ---------------------------------------------------------------------
# Public alias (short name used in examples)
# ---------------------------------------------------------------------


# ---------------------------------------------------------------------
# Public alias (short name used in examples)
# ---------------------------------------------------------------------
SandevistanTF = SandevistanKerasCallback
