import unittest

import torch

from backbones import TemporalAttention


class TemporalAttentionTest(unittest.TestCase):
    def test_zero_rbf_bias_has_live_projection_gradient(self) -> None:
        torch.manual_seed(0)
        attention = TemporalAttention(
            n_embd=12,
            n_head=3,
            use_time_rope=False,
            use_rbf_bias=True,
        )
        features = torch.randn(2, 4, 4, 16)
        target = torch.randn(2, 4, 4, 3)

        initial_bias = (
            attention.time_bias_scale.tanh()
            * attention.rbf_proj(features)
        )
        torch.testing.assert_close(initial_bias, torch.zeros_like(initial_bias))

        (initial_bias * target).sum().backward()
        projection_grad = attention.rbf_proj.weight.grad
        self.assertIsNotNone(projection_grad)
        self.assertGreater(projection_grad.abs().sum().item(), 0.0)

        with torch.no_grad():
            attention.rbf_proj.weight.add_(projection_grad, alpha=-1e-3)
        attention.zero_grad(set_to_none=True)
        updated_bias = (
            attention.time_bias_scale.tanh()
            * attention.rbf_proj(features)
        )
        (updated_bias * target).sum().backward()

        scale_grad = attention.time_bias_scale.grad
        self.assertIsNotNone(scale_grad)
        self.assertGreater(scale_grad.abs().item(), 0.0)


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