import math
import unittest

from dropout_decay.schedules import DropoutDecayConfig, DropoutDecayScheduler
from dropout_decay.specs import DropoutCondition, anchor_dropout


class DropoutScheduleTests(unittest.TestCase):
    def test_linear_decay_endpoints_and_midpoint(self):
        scheduler = DropoutDecayScheduler(
            DropoutDecayConfig(
                initial_dropout=0.30,
                final_dropout=0.10,
                decay_tokens=100,
                schedule="linear",
            )
        )

        self.assertAlmostEqual(scheduler.value(0), 0.30)
        self.assertAlmostEqual(scheduler.value(50), 0.20)
        self.assertAlmostEqual(scheduler.value(100), 0.10)
        self.assertAlmostEqual(scheduler.value(1000), 0.10)

    def test_cosine_decay_midpoint(self):
        scheduler = DropoutDecayScheduler(
            DropoutDecayConfig(
                initial_dropout=0.40,
                final_dropout=0.20,
                decay_tokens=100,
                schedule="cosine",
            )
        )

        self.assertAlmostEqual(scheduler.value(50), 0.30)

    def test_invalid_scheduler_config_rejected(self):
        with self.assertRaises(ValueError):
            DropoutDecayScheduler(
                DropoutDecayConfig(
                    initial_dropout=0.10,
                    final_dropout=0.30,
                    decay_tokens=100,
                    schedule="linear",
                )
            )

    def test_anchor_dropout_interpolates_in_log_token_space(self):
        anchors = ((100, 0.30), (10_000, 0.10))
        value = anchor_dropout(1_000, anchors)
        expected_mix = (math.log(1_000) - math.log(100)) / (
            math.log(10_000) - math.log(100)
        )
        expected = 0.30 + expected_mix * (0.10 - 0.30)

        self.assertAlmostEqual(value, expected)

    def test_anchor_condition_requires_unique_tokens(self):
        condition = DropoutCondition(
            name="fit",
            kind="anchor_decay",
            initial=0.30,
            final=0.10,
            anchors=((100, 0.30), (10_000, 0.10)),
        )

        with self.assertRaises(ValueError):
            condition.make_fn(fallback_decay_tokens=100)

        fn = condition.make_fn(fallback_decay_tokens=100, unique_tokens=1_000)
        self.assertAlmostEqual(fn(0), anchor_dropout(1_000, condition.anchors))


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