import unittest
from beamboot.tinylm import TinyLM, find_weights, reference, LANES, ITERS

CYCLE = [(c + 1) % LANES for c in range(LANES)]


class TinyLMTests(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.weights = find_weights(CYCLE)
        cls.lm = TinyLM(cls.weights)

    def test_reference_realizes_cycle(self):
        self.assertEqual(reference(self.weights, 0, 8), [1, 2, 3, 4, 5, 6, 7, 0])

    def test_generate_runs_the_network_inside_the_model(self):
        chars, tokens = self.lm.generate(0, ITERS)
        self.assertEqual(len(tokens), 42*ITERS)
        self.assertEqual(chars, reference(self.weights, 0, ITERS))
        self.assertEqual(chars, [1, 2, 3, 4, 5, 6, 7, 0, 1, 2])

    def test_other_start_char(self):
        chars, _ = self.lm.generate(5, 4)
        self.assertEqual(chars, [6, 7, 0, 1])


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