import unittest

import torch

from predictor_training.confidence import ConfidenceTokenHead


class ConfidenceTokenHeadTest(unittest.TestCase):
    def setUp(self) -> None:
        torch.manual_seed(0)
        self.head = ConfidenceTokenHead(
            dim=16,
            token_dim=8,
            context_dim=4,
            num_heads=2,
            ffn_dim=32,
            dropout=0.0,
            num_steps=3,
        )

    def test_output_shape_and_predictor_features_are_detached(self) -> None:
        transformed = torch.randn(2, 7, 16, requires_grad=True)
        predicted = torch.randn(2, 7, 16, requires_grad=True)
        anchor = torch.randn(2, 7, 16, requires_grad=True)
        output = self.head(
            transformed_hidden=transformed,
            pred_hidden=predicted,
            anchor_hidden=anchor,
            chunk_position=torch.tensor([0.0, 1.0]),
            step_id=torch.tensor([1, 3]),
        )
        self.assertEqual(output.shape, (2,))
        output.sum().backward()
        self.assertIsNone(transformed.grad)
        self.assertIsNone(predicted.grad)
        self.assertIsNone(anchor.grad)
        self.assertTrue(any(parameter.grad is not None for parameter in self.head.parameters()))

    def test_rejects_unsupported_step(self) -> None:
        value = torch.randn(1, 7, 16)
        with self.assertRaisesRegex(ValueError, "supports step_id"):
            self.head(
                transformed_hidden=value,
                pred_hidden=value,
                anchor_hidden=value,
                chunk_position=torch.tensor([0.0]),
                step_id=torch.tensor([4]),
            )

    def test_production_dimensions_and_parameter_count(self) -> None:
        head = ConfidenceTokenHead()
        self.assertEqual(head.token_projection[0].in_features, 1536)
        self.assertEqual(head.token_projection[0].out_features, 512)
        self.assertEqual(head.output[0].in_features, 1156)
        count = sum(parameter.numel() for parameter in head.parameters())
        self.assertEqual(count, 4_872_065)


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