# Path: sockeye/lr_scheduler.py
# class LearningRateScheduler:
# class AdaptiveLearningRateScheduler(LearningRateScheduler):
# class LearningRateSchedulerInvSqrtDecay(LearningRateScheduler):
# class LearningRateSchedulerLinearDecay(LearningRateScheduler):
# class LearningRateSchedulerPlateauReduce(AdaptiveLearningRateScheduler):
#     def __init__(self, base_lr: float = 1.0, warmup: int = 0, t_scale: float = 1.0) -> None:
#     def __call__(self, t):
#     def _warmup(self, scaled_t):
#     def new_evaluation_result(self, has_improved: bool) -> bool:
#     def __call__(self, t: int):
#     def __init__(self, base_lr: float, total_steps: int, warmup: int = 0, t_scale: float = 1.0) -> None:
#     def __call__(self, t: int):
#     def __init__(self, base_lr: float, reduce_factor: float, reduce_num_not_improved: int, warmup: int = 0) -> None:
#     def new_evaluation_result(self, has_improved: bool) -> bool:
#     def __call__(self, t):
#     def __repr__(self):
# def get_lr_scheduler(scheduler_type: str,
#                      base_learning_rate: float,
#                      learning_rate_t_scale: float,
#                      learning_rate_reduce_factor: float,
#                      learning_rate_reduce_num_not_improved: int,
#                      learning_rate_warmup: int = 0,
#                      max_updates: Optional[int] = None) -> Optional[LearningRateScheduler]:


import pytest
import numpy as np
from sockeye import lr_scheduler
                                              base_learning_rate=1,
                                              learning_rate_t_scale=1,
                                              learning_rate_reduce_factor=0,
                                              learning_rate_reduce_num_not_improved=0,
                                              learning_rate_warmup=3,
                                              max_updates=10)

    # Warmup term * decay term
    expected_schedule = [
        (1/3) * (9/10),
        (2/3) * (8/10),
        (3/3) * (7/10),
        (3/3) * (6/10),
        (3/3) * (5/10),
        (3/3) * (4/10),
        (3/3) * (3/10),
        (3/3) * (2/10),
        (3/3) * (1/10),
        (3/3) * (0/10),
    ]
    actual_schedule = [scheduler(t) for t in range(1, 11)]
    assert np.isclose(expected_schedule, actual_schedule).all()


@pytest.mark.parametrize('scheduler_type, expected_instance',
                         [('none', None),
                          ('inv-sqrt-decay', lr_scheduler.LearningRateSchedulerInvSqrtDecay),
                          ('linear-decay', lr_scheduler.LearningRateSchedulerLinearDecay),
                          ('plateau-reduce', lr_scheduler.LearningRateSchedulerPlateauReduce)])
def test_get_lr_scheduler(scheduler_type, expected_instance):

# Path: sockeye/lr_scheduler.py
# class LearningRateScheduler:
# class AdaptiveLearningRateScheduler(LearningRateScheduler):
# class LearningRateSchedulerInvSqrtDecay(LearningRateScheduler):
# class LearningRateSchedulerLinearDecay(LearningRateScheduler):
# class LearningRateSchedulerPlateauReduce(AdaptiveLearningRateScheduler):
#     def __init__(self, base_lr: float = 1.0, warmup: int = 0, t_scale: float = 1.0) -> None:
#     def __call__(self, t):
#     def _warmup(self, scaled_t):
#     def new_evaluation_result(self, has_improved: bool) -> bool:
#     def __call__(self, t: int):
#     def __init__(self, base_lr: float, total_steps: int, warmup: int = 0, t_scale: float = 1.0) -> None:
#     def __call__(self, t: int):
#     def __init__(self, base_lr: float, reduce_factor: float, reduce_num_not_improved: int, warmup: int = 0) -> None:
#     def new_evaluation_result(self, has_improved: bool) -> bool:
#     def __call__(self, t):
#     def __repr__(self):
# def get_lr_scheduler(scheduler_type: str,
#                      base_learning_rate: float,
#                      learning_rate_t_scale: float,
#                      learning_rate_reduce_factor: float,
#                      learning_rate_reduce_num_not_improved: int,
#                      learning_rate_warmup: int = 0,
#                      max_updates: Optional[int] = None) -> Optional[LearningRateScheduler]:


# Path: test/unit/test_lr_scheduler.py
import pytest
import numpy as np
from sockeye import lr_scheduler

                                              base_learning_rate=1,
                                              learning_rate_t_scale=1,
                                              learning_rate_reduce_factor=0,
                                              learning_rate_reduce_num_not_improved=0,
                                              learning_rate_warmup=3,
                                              max_updates=10)

    # Warmup term * decay term
    expected_schedule = [
        (1/3) * (9/10),
        (2/3) * (8/10),
        (3/3) * (7/10),
        (3/3) * (6/10),
        (3/3) * (5/10),
        (3/3) * (4/10),
        (3/3) * (3/10),
        (3/3) * (2/10),
        (3/3) * (1/10),
        (3/3) * (0/10),
    ]
    actual_schedule = [scheduler(t) for t in range(1, 11)]
    assert np.isclose(expected_schedule, actual_schedule).all()


@pytest.mark.parametrize('scheduler_type, expected_instance',
                         [('none', None),
                          ('inv-sqrt-decay', lr_scheduler.LearningRateSchedulerInvSqrtDecay),
                          ('linear-decay', lr_scheduler.LearningRateSchedulerLinearDecay),
                          ('plateau-reduce', lr_scheduler.LearningRateSchedulerPlateauReduce)])
def test_get_lr_scheduler(scheduler_type, expected_instance):

    scheduler = lr_scheduler.get_lr_scheduler(scheduler_type,