import unittest

import numpy as np
import torch

from wannabe_tokenizer.models.pipeline import chunk_ranges
from wannabe_tokenizer.quantizer import CVQHead, to_token_clock


class TokenClockTests(unittest.TestCase):
    def test_offset_delays_and_repeats_edges(self):
        values = np.arange(5)
        np.testing.assert_array_equal(to_token_clock(values, 1), [0, 0, 1, 2, 3])
        np.testing.assert_array_equal(to_token_clock(values, 0), values)

    def test_context_chunks_cover_every_frame_once(self):
        ranges = chunk_ranges(100, prefix_tokens=24, context=100)
        self.assertEqual(ranges[0][0], 0)
        self.assertEqual(ranges[-1][1], 100)
        self.assertEqual(sum(end - start for start, end in ranges), 100)
        for (_, previous_end), (next_start, _) in zip(ranges, ranges[1:]):
            self.assertEqual(previous_end, next_start)

    def test_cvq_head_output_shape(self):
        torch.manual_seed(0)
        head = CVQHead(width=16, dim=4, hidden=8, tau=0.07).eval()
        output = head(torch.randn(2, 3, 16))
        self.assertEqual(tuple(output.shape), (2, 3, 32768))
        self.assertTrue(torch.isfinite(output).all())


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