from typing import Union
import tensorflow as tf

from sandevistanlr.core.schedule_base import SandevistanScheduleBase


class SandevistanLRTF:
    """
    TensorFlow adapter for SandevistanScheduleBase.

    Converts a global step (int or tf.Tensor) into
    a TensorFlow scalar learning rate.
    """

    def __init__(self, schedule: SandevistanScheduleBase):
        self.schedule = schedule

    def __call__(self, step: Union[int, tf.Tensor]) -> tf.Tensor:
        if isinstance(step, tf.Tensor):
            step = int(step.numpy().reshape(-1)[0])
        else:
            step = int(step)

        lr = self.schedule(step)
        return tf.convert_to_tensor(lr, dtype=tf.float32)
