import random
import unittest
import torch
from beamboot.machine import NeuralMachine, State, decode_output, reference_step, run, state_bits


class MachineTests(unittest.TestCase):
    def test_all_states(self):
        program = [('LDI', 7), ('STM', 0), ('SUB', 1), ('JZ', 6), ('ADD', 2), ('JMP', 2), ('LDM', 0), ('OUT', 0), ('HALT', 0)]
        model = NeuralMachine(program)
        states = [State(pc, acc, ram, out) for pc in range(16) for acc in range(16) for ram in range(16) for out in (False, True)]
        with torch.inference_mode():
            for start in range(0, len(states), 256):
                chunk = states[start:start+256]
                actual = model(torch.stack([state_bits(s) for s in chunk]))
                self.assertEqual([decode_output(row) for row in actual], [reference_step(program, s) for s in chunk])

    def test_writable_memory(self):
        program = [('LDI', 11), ('STM', 0), ('LDI', 0), ('LDM', 0), ('OUT', 0), ('HALT', 0)]
        states = list(run(NeuralMachine(program)))
        self.assertEqual([(s.acc, s.ram) for s in states if s.out], [(11, 11)])
        self.assertTrue(states[-1].halt)