import unittest
import torch
from beamboot.sequence_cpu import SequenceCPU, SequenceRuntime, event
from beamboot.monitor import Monitor, firmware


class SequenceTests(unittest.TestCase):
    def test_full_attention_is_causal_and_last_decode_matches(self):
        model = SequenceCPU([('LOAD', 2), ('HALT', 0)])
        events = [event(address=i, valid=True) for i in range(16)]
        events += [event(address=2, value=7, valid=True, position=16), event(position=17)]
        x = torch.tensor([events], dtype=torch.float32)
        with torch.inference_mode():
            first = model(x)
            cached = model(x, last_only=True)
            changed = x.clone()
            changed[0, -1, 14:18] = 1
            changed_result = model(changed)
        self.assertTrue(torch.equal(first[:, -1:], cached))
        self.assertTrue(torch.equal(first[:, :-1], changed_result[:, :-1]))
        self.assertEqual(first[0, -1, 5:9].tolist(), [1, 1, 1, 0])

    def test_stores_and_reloads(self):
        program = [('LDI', 9), ('STORE', 4), ('LDI', 2), ('STORE', 5),
                   ('LDI', 0), ('LOAD', 4), ('OUT', 0), ('LOAD', 5), ('OUT', 0), ('HALT', 0)]
        runtime = SequenceRuntime(SequenceCPU(program))
        steps = []
        while not runtime.halted: steps.append(runtime.step())
        self.assertEqual([s.acc for s in steps if s.output], [9, 2])
        self.assertEqual([(s.address, s.value) for s in steps if s.write], [(4, 9), (5, 2)])

    def test_monitor_firmware(self):
        monitor = Monitor()
        monitor.runtime = SequenceRuntime(SequenceCPU(firmware()))
        self.assertEqual(monitor.boot(), b'>')
        self.assertEqual(monitor.command(1), b'OK')
        self.assertEqual(monitor.command(2), b'HI')
        self.assertEqual(monitor.command(3), b'')
        self.assertEqual(monitor.command(1), b'OK')

    def test_bounded_context(self):
        machine = SequenceRuntime(SequenceCPU([('JMP', 0)], max_events=18))
        machine.step(); machine.step()
        with self.assertRaisesRegex(RuntimeError, 'budget'): machine.step()