from typing import Any, List

from transformers import (
    get_constant_schedule_with_warmup,
    get_cosine_schedule_with_warmup,
    get_linear_schedule_with_warmup,
)

__all__ = ["Schedulers"]


def constant_schedule_with_warmup(optimizer, num_warmup_steps, **kwargs):
    return get_constant_schedule_with_warmup(
        optimizer=optimizer, num_warmup_steps=num_warmup_steps
    )


class Schedulers:
    """Schedulers factory."""

    _schedulers = {
        "Cosine": get_cosine_schedule_with_warmup,
        "Linear": get_linear_schedule_with_warmup,
        "Constant": constant_schedule_with_warmup,
    }

    @classmethod
    def names(cls) -> List[str]:
        return sorted(cls._schedulers.keys())

    @classmethod
    def get(cls, name: str) -> Any:
        """Access to Schedulers.

        Args:
            name: scheduler name
        Returns:
            A class to build the Schedulers
        """
        return cls._schedulers.get(name)