from __future__ import annotations

from typing import Any, Dict
import tensorflow as tf

from .schedule_base import SandevistanScheduleBase


class SandevistanLRTF:
    """
    TensorFlow-compatible wrapper for SandevistanScheduleBase.

    Intentionally does NOT subclass LearningRateSchedule
    because TensorFlow typing stubs mark it as Any.
    """

    def __init__(self, base: SandevistanScheduleBase) -> None:
        self.base = base

    def __call__(self, step: tf.Tensor) -> tf.Tensor:
        step_i = tf.cast(step, tf.int64)

        def _lr_fn(s: tf.Tensor) -> tf.Tensor:
            lr = self.base.lr_at_step(int(s.numpy()))
            return tf.constant(lr, dtype=tf.float32)

        return tf.map_fn(
            _lr_fn,
            step_i,
            fn_output_signature=tf.float32,
        )

    def get_config(self) -> Dict[str, Any]:
        return {
            "steps_per_epoch": self.base.steps_per_epoch,
            "lr_min": self.base.lr_min,
            "lr_max": self.base.lr_max,
            "epoch_plan": self.base.epoch_plan,
        }
