import pytest
import torch
import torch.nn as nn
from torch_surgeon import Surgeon


def simple_model():
    return nn.Sequential(
        nn.Linear(32, 16),
        nn.ReLU(),
        nn.Linear(16, 8),
        nn.ReLU(),
        nn.Linear(8, 1),
    )


def run_steps(model, surgeon, n=5):
    opt = torch.optim.Adam(model.parameters(), lr=1e-3)
    crit = nn.MSELoss()
    for _ in range(n):
        x = torch.randn(8, 32)
        y = torch.randn(8, 1)
        loss = crit(model(x), y)
        opt.zero_grad()
        loss.backward()
        opt.step()


class TestAttachDetach:
    def test_attach_registers_hooks(self):
        model = simple_model()
        surgeon = Surgeon(model, rules='default')
        assert not surgeon.is_attached
        surgeon.attach()
        assert surgeon.is_attached
        assert len(surgeon._hooks) > 0

    def test_detach_removes_all_hooks(self):
        model = simple_model()
        surgeon = Surgeon(model).attach()
        surgeon.detach()
        assert not surgeon.is_attached
        assert len(surgeon._hooks) == 0

    def test_double_attach_is_safe(self):
        model = simple_model()
        surgeon = Surgeon(model).attach()
        hooks_before = len(surgeon._hooks)
        surgeon.attach()  # should be a no-op
        assert len(surgeon._hooks) == hooks_before

    def test_step_counter_increments(self):
        model = simple_model()
        surgeon = Surgeon(model).attach()
        run_steps(model, surgeon, n=3)
        assert surgeon.step == 3


class TestReport:
    def test_report_returns_dict(self):
        model = simple_model()
        surgeon = Surgeon(model).attach()
        run_steps(model, surgeon, n=2)
        report = surgeon.report()
        assert isinstance(report, dict)

    def test_report_contains_layer_keys(self):
        model = simple_model()
        surgeon = Surgeon(model).attach()
        run_steps(model, surgeon, n=2)
        report = surgeon.report()
        # Should have entries for leaf layers
        assert len(report) > 0

    def test_report_value_structure(self):
        model = simple_model()
        surgeon = Surgeon(model).attach()
        run_steps(model, surgeon, n=2)
        for name, stats in surgeon.report().items():
            assert 'mean' in stats
            assert 'std' in stats
            assert 'norm' in stats
            assert 'status' in stats
            assert stats['status'] in ('ok', 'vanishing', 'exploding', 'stagnant')


class TestReset:
    def test_reset_clears_stats(self):
        model = simple_model()
        surgeon = Surgeon(model).attach()
        run_steps(model, surgeon, n=3)
        surgeon.reset()
        assert surgeon.step == 0
        assert surgeon.report() == {}


class TestCustomRules:
    def test_custom_rules_dict(self):
        model = simple_model()
        surgeon = Surgeon(model, rules={
            'vanishing_threshold': 0.05,
            'exploding_threshold': 50.0,
            'log_every': 5,
            'verbose': False,
        }).attach()
        run_steps(model, surgeon, n=2)
        assert surgeon.is_attached

    def test_invalid_rules_raises(self):
        model = simple_model()
        with pytest.raises(ValueError):
            Surgeon(model, rules='invalid')