import unittest

import torch

from backbones import TrajMixer


class TrajMixerTest(unittest.TestCase):
    def test_default_shape_parameters_and_initialization(self) -> None:
        mixer = TrajMixer(
            n_embd=120,
            n_head=10,
            dropout=0.0,
        )

        x = torch.randn(2, 7, 120)
        self.assertEqual(mixer(x).shape, x.shape)
        self.assertEqual(sum(p.numel() for p in mixer.parameters()), 32_040)
        self.assertEqual(tuple(mixer.norm.normalized_shape), (120,))
        self.assertEqual(tuple(mixer.intra_gate_logits.shape), (10, 12))
        torch.testing.assert_close(
            torch.sigmoid(mixer.intra_gate_logits.detach()),
            torch.full((10, 12), 0.1),
        )
        self.assertEqual(mixer.intra_hidden, 48)
        self.assertEqual(
            tuple(mixer.intra_gate_proj.shape),
            (10, 12, 48),
        )
        self.assertEqual(
            tuple(mixer.intra_value_proj.shape),
            (10, 12, 48),
        )
        self.assertEqual(
            tuple(mixer.intra_output_proj.shape),
            (10, 48, 12),
        )
        self.assertEqual(mixer.hidden_group, 40)
        self.assertEqual(tuple(mixer.gate_proj.shape), (12, 10, 40))
        self.assertEqual(tuple(mixer.value_proj.shape), (12, 10, 40))
        self.assertEqual(tuple(mixer.output_proj.shape), (12, 40, 10))

    def test_zero_final_output_projection_makes_mixer_identity(self) -> None:
        torch.manual_seed(0)
        mixer = TrajMixer(12, n_head=3, dropout=0.0)
        with torch.no_grad():
            mixer.output_proj.zero_()
        x = torch.randn(2, 5, 12)
        torch.testing.assert_close(mixer(x), x)

    def test_forward_matches_single_outer_residual_formula(self) -> None:
        torch.manual_seed(0)
        mixer = TrajMixer(12, n_head=3, dropout=0.0)
        mixer.eval()
        x = torch.randn(2, 5, 12)

        grouped = mixer.norm(x).reshape(2, 5, 3, 4)
        intra_output = mixer._intra_mix(grouped)
        static_gate = torch.sigmoid(mixer.intra_gate_logits).view(
            1, 1, 3, 4
        )
        mixed_input = grouped + static_gate * intra_output
        update = mixer._cross_mix(mixed_input).reshape(2, 5, 12)

        torch.testing.assert_close(mixer(x), x + update)

    def test_intra_stage_is_independent_across_groups(self) -> None:
        torch.manual_seed(0)
        mixer = TrajMixer(12, n_head=3, dropout=0.0)
        mixer.eval()

        grouped = torch.randn(2, 4, 3, 4)
        changed = grouped.clone()
        changed[:, :, 1, :] += torch.randn_like(changed[:, :, 1, :])

        original_out = mixer._intra_mix(grouped)
        changed_out = mixer._intra_mix(changed)
        unchanged_groups = torch.tensor([0, 2])
        torch.testing.assert_close(
            original_out.index_select(2, unchanged_groups),
            changed_out.index_select(2, unchanged_groups),
        )

    def test_cross_stage_mixes_groups_without_mixing_coordinates(self) -> None:
        mixer = TrajMixer(6, n_head=3, dropout=0.0)
        mixer.eval()
        with torch.no_grad():
            mixer.gate_proj.zero_()
            mixer.value_proj.zero_()
            mixer.output_proj.zero_()

            # Coordinate 0 reads group 0 through hidden unit 0 and writes it
            # into group 1. Coordinate 1 must remain independent.
            mixer.gate_proj[0, 0, 0] = 1.0
            mixer.value_proj[0, 0, 0] = 1.0
            mixer.output_proj[0, 0, 1] = 1.0

        grouped = torch.tensor(
            [[[
                [-1.0, 4.0],
                [0.0, 5.0],
                [1.0, 6.0],
            ]]]
        )
        changed = grouped.clone()
        changed[0, 0, 0, 0] = 2.0

        original_out = mixer._cross_mix(grouped)
        changed_out = mixer._cross_mix(changed)

        self.assertNotEqual(
            original_out[0, 0, 1, 0].item(),
            changed_out[0, 0, 1, 0].item(),
        )
        torch.testing.assert_close(
            original_out[..., 1],
            changed_out[..., 1],
        )

    def test_mixer_does_not_mix_sequence_positions(self) -> None:
        torch.manual_seed(0)
        mixer = TrajMixer(12, n_head=3, dropout=0.0)
        mixer.eval()
        x = torch.randn(2, 5, 12)
        changed = x.clone()
        changed[:, 3, :] += torch.randn_like(changed[:, 3, :])

        original_out = mixer(x)
        changed_out = mixer(changed)
        unchanged_positions = torch.tensor([0, 1, 2, 4])
        torch.testing.assert_close(
            original_out.index_select(1, unchanged_positions),
            changed_out.index_select(1, unchanged_positions),
        )

    def test_gradients_reach_every_projection_family(self) -> None:
        torch.manual_seed(1)
        mixer = TrajMixer(12, n_head=3, dropout=0.0)
        x = torch.randn(2, 4, 12, requires_grad=True)

        mixer(x).square().mean().backward()

        self.assertIsNotNone(x.grad)
        self.assertTrue(torch.isfinite(x.grad).all())
        self.assertGreater(x.grad.abs().sum().item(), 0.0)
        projection_names = (
            "intra_gate_proj",
            "intra_value_proj",
            "intra_output_proj",
            "gate_proj",
            "value_proj",
            "output_proj",
        )
        for name in projection_names:
            parameter = getattr(mixer, name)
            self.assertIsNotNone(parameter.grad, name)
            self.assertTrue(torch.isfinite(parameter.grad).all(), name)
            self.assertGreater(parameter.grad.abs().sum().item(), 0.0, name)

    def test_invalid_group_partition_is_rejected(self) -> None:
        with self.assertRaisesRegex(ValueError, "divisible"):
            TrajMixer(n_embd=121, n_head=10)


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