import pytest
from torch_surgeon.detector import GradientDetector


class TestVanishingDetection:
    def test_detects_vanishing(self):
        det = GradientDetector(vanishing_threshold=0.1, ema_alpha=1.0)
        # Prime the EMA with a large norm
        det.update('layer1', 0.5, 0.1, 1.0)
        # Now send a very small norm — ratio = 0.0001 / 1.0 = 0.0001 < 0.1
        status = det.update('layer1', 0.0, 0.0, 0.0001)
        assert status == 'vanishing'

    def test_ok_when_normal(self):
        det = GradientDetector(ema_alpha=1.0)
        det.update('layer1', 0.5, 0.1, 1.0)
        status = det.update('layer1', 0.5, 0.1, 0.9)
        assert status == 'ok'


class TestExplodingDetection:
    def test_detects_exploding(self):
        det = GradientDetector(exploding_threshold=10.0, ema_alpha=1.0)
        det.update('layer1', 0.5, 0.1, 1.0)
        # ratio = 500 / 1.0 = 500 > 10.0
        status = det.update('layer1', 50.0, 10.0, 500.0)
        assert status == 'exploding'


class TestStagnantDetection:
    def test_detects_stagnant(self):
        det = GradientDetector(stagnant_steps=3)
        for _ in range(3):
            status = det.update('layer1', 0.0, 0.0, 1e-10)
        assert status == 'stagnant'

    def test_stagnant_resets_on_activity(self):
        det = GradientDetector(stagnant_steps=3, ema_alpha=0.5)
        for _ in range(2):
            det.update('layer1', 0.0, 0.0, 1e-10)
        # Activity resets counter
        status = det.update('layer1', 0.5, 0.1, 1.0)
        assert det._stagnant_count.get('layer1', 0) == 0


class TestGetStats:
    def test_get_stats_returns_copy(self):
        det = GradientDetector()
        det.update('a', 0.1, 0.01, 1.0)
        s1 = det.get_stats()
        s2 = det.get_stats()
        assert s1 == s2
        assert s1 is not s2

    def test_reset_clears_all(self):
        det = GradientDetector()
        det.update('a', 0.1, 0.01, 1.0)
        det.reset()
        assert det.get_stats() == {}