import unittest
import numpy as np
from attention import rope_2d, scaled_dot_product_attention


class AttentionTests(unittest.TestCase):
    def test_attention_probabilities_sum_to_one(self):
        x = np.array([[1.0, 0.0], [0.0, 1.0]])
        output, weights = scaled_dot_product_attention(x, x, x)
        self.assertEqual(output.shape, (2, 2))
        np.testing.assert_allclose(weights.sum(axis=-1), np.ones(2))

    def test_causal_attention_blocks_future_tokens(self):
        x = np.eye(3)
        _, weights = scaled_dot_product_attention(x, x, x, causal=True)
        np.testing.assert_allclose(np.triu(weights, k=1), np.zeros((3, 3)))

    def test_rope_preserves_vector_norm(self):
        values = np.array([[3.0, 4.0], [1.0, -2.0]])
        rotated = rope_2d(values, np.array([0.0, 1.0]))
        np.testing.assert_allclose(np.linalg.norm(rotated, axis=-1), np.linalg.norm(values, axis=-1))


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