from types import SimpleNamespace
import pytest
from trainer.distillation import _cosine_learning_rate, _learning_rate_schedule_identity


def test_step80_keeps_the_original_cosine_curve():
    cfg = SimpleNamespace(lr_schedule="cosine", lr=2.5e-6, lr_final=5e-7,
                          lr_schedule_total_steps=200, max_steps=80)
    early = _learning_rate_schedule_identity(cfg)
    cfg.max_steps = 200
    original = _learning_rate_schedule_identity(cfg)
    assert early == original
    for update in range(80):
        assert _cosine_learning_rate(update, start=early["start"], end=early["end"], total_updates=early["total_updates"]) == _cosine_learning_rate(update, start=2.5e-6, end=5e-7, total_updates=200)
    assert _cosine_learning_rate(79, start=2.5e-6, end=5e-7, total_updates=200) > 5e-7
    cfg.max_steps = 1
    assert _learning_rate_schedule_identity(cfg) == original
    cfg.max_steps = 201
    with pytest.raises(ValueError):
        _learning_rate_schedule_identity(cfg)