import unittest

from scripts.train_confidence_token_lazy_ddp import learning_rate_at_step


class ConfidenceLearningRateTest(unittest.TestCase):
    def value(self, step: int) -> float:
        return learning_rate_at_step(
            step=step,
            total_steps=1440,
            schedule="warmup_cosine",
            peak_lr=1e-4,
            warmup_steps=36,
            warmup_start_lr=1e-5,
            min_lr=1e-6,
        )

    def test_warmup_endpoints(self) -> None:
        self.assertAlmostEqual(self.value(0), 1e-5)
        self.assertAlmostEqual(self.value(35), 1e-4)

    def test_cosine_is_monotonic_and_reaches_floor(self) -> None:
        values = [self.value(step) for step in range(35, 1440)]
        self.assertTrue(all(left >= right for left, right in zip(values, values[1:])))
        self.assertAlmostEqual(values[-1], 1e-6)

    def test_constant_schedule_is_unchanged(self) -> None:
        value = learning_rate_at_step(
            step=503,
            total_steps=1440,
            schedule="constant",
            peak_lr=3e-4,
            warmup_steps=0,
            warmup_start_lr=1e-5,
            min_lr=1e-6,
        )
        self.assertEqual(value, 3e-4)


if __name__ == "__main__":
    unittest.main()