import unittest

import torch

from speech_bridge_gemma.longcat_semantic_stride_eval import compact_semantic, expand_semantic, stride_semantic


class LongCatSemanticStrideEvalTest(unittest.TestCase):
    def test_compact_first_keeps_first_code_per_chunk(self) -> None:
        out = compact_semantic(torch.tensor([1, 2, 3, 4, 5]), 2, "first")

        self.assertEqual(out.tolist(), [1, 3, 5])

    def test_compact_center_keeps_center_code_per_chunk(self) -> None:
        out = compact_semantic(torch.tensor([1, 2, 3, 4, 5]), 2, "center")

        self.assertEqual(out.tolist(), [2, 4, 5])

    def test_compact_mode_keeps_most_common_code_per_chunk(self) -> None:
        out = compact_semantic(torch.tensor([1, 2, 2, 4, 5]), 3, "mode")

        self.assertEqual(out.tolist(), [2, 4])

    def test_expand_repeats_and_truncates_to_frame_count(self) -> None:
        out = expand_semantic(torch.tensor([7, 8, 9]), 5, 2)

        self.assertEqual(out.tolist(), [7, 7, 8, 8, 9])

    def test_stride_semantic_returns_compact_and_expanded(self) -> None:
        compact, expanded = stride_semantic(torch.tensor([1, 2, 3, 4]), 4, 2, "first")

        self.assertEqual(compact.tolist(), [1, 3])
        self.assertEqual(expanded.tolist(), [1, 1, 3, 3])


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