import unittest
import torch
from beamboot.cpu import CompiledCPU, Instruction, HALT, VOCAB, encode, output_bytes, print_program, run, transition


class CPUTests(unittest.TestCase):
    def test_all_states_match_reference(self):
        program = [Instruction('LDI', 3), Instruction('SUB', 1), Instruction('JZ', 4), Instruction('JMP', 1), Instruction('OUT'), Instruction('HALT')]
        model = CompiledCPU(program)
        with torch.inference_mode():
            result = model(torch.arange(VOCAB)).argmax(dim=-1).tolist()
        self.assertEqual(result, [transition(program, i) for i in range(VOCAB)])

    def test_boot_text(self):
        model = CompiledCPU(print_program('OK\n'))
        self.assertEqual(output_bytes(list(run(model))), b'OK\n')

    def test_loop_and_output(self):
        program = [Instruction('LDI', 3), Instruction('OUT'), Instruction('SUB', 1), Instruction('JZ', 5), Instruction('JMP', 1), Instruction('HALT')]
        result = list(run(CompiledCPU(program)))
        self.assertEqual(result[-1], HALT)
        self.assertEqual([t // 2 % 16 for t in result if t != HALT and t % 2], [3, 2, 1])

    def test_infinite_loop_budget(self):
        with self.assertRaises(RuntimeError):
            list(run(CompiledCPU([Instruction('JMP', 0)]), 3))