import random
import unittest
from beamboot.vector_lm import VectorLM, LANES, ROWS, GENS


class VectorLMTests(unittest.TestCase):
    def test_model_emits_exact_dot_as_next_token(self):
        rng = random.Random(417)
        weights = [[rng.randint(-8, 7) for _ in range(LANES)] for _ in range(ROWS)]
        machine = VectorLM(weights)
        for _ in range(GENS):
            row = rng.randrange(ROWS)
            acts = [rng.randint(0, 15) for _ in range(LANES)]
            expected = sum(w * a for w, a in zip(weights[row], acts))
            self.assertEqual(machine.dot(row, acts), expected)

    def test_extremes(self):
        weights = [[-8]*LANES, [7]*LANES] + [[0]*LANES]*(ROWS-2)
        machine = VectorLM(weights)
        self.assertEqual(machine.dot(0, [15]*LANES), -960)
        self.assertEqual(machine.dot(1, [15]*LANES), 840)


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