import random
import unittest
import torch
from beamboot.large_cpu import LargeCPU, LargeRuntime, encode_event


class LargeCPUEquivalenceTests(unittest.TestCase):
    def test_random_programs_against_independent_interpreter(self):
        rng = random.Random(2401)
        ops = ('LDI', 'ADD', 'SUB', 'LOAD', 'STORE', 'ADDM', 'SUBM', 'OUT')
        for _ in range(4):
            program = [(rng.choice(ops), rng.randrange(16)) for _ in range(70)] + [('HALT', 0)]
            machine = LargeRuntime(LargeCPU(program))
            memory, accumulator = [0]*16, 0
            for pc, (op, arg) in enumerate(program):
                step = machine.step()
                if op == 'LDI': accumulator = arg
                elif op == 'ADD': accumulator = (accumulator + arg) & 15
                elif op == 'SUB': accumulator = (accumulator - arg) & 15
                elif op == 'LOAD': accumulator = memory[arg]
                elif op == 'STORE': memory[arg] = accumulator
                elif op == 'ADDM': accumulator = (accumulator + memory[arg]) & 15
                elif op == 'SUBM': accumulator = (accumulator - memory[arg]) & 15
                self.assertEqual(step.acc, accumulator)
                self.assertEqual(step.pc, pc if op == 'HALT' else pc + 1)
                self.assertEqual(step.write, op == 'STORE')
                self.assertEqual(step.output, op == 'OUT')
                self.assertEqual(step.halt, op == 'HALT')
                if step.write: self.assertEqual((step.address, step.value), (arg, accumulator))

    def test_branches_and_high_rom_address(self):
        program = [('LDI', 3), ('SUB', 1), ('JZ', 40), ('JMP', 1)]
        program += [('HALT', 0)] * (40-len(program))
        program += [('LDI', 9), ('OUT', 0), ('HALT', 0)]
        runtime = LargeRuntime(LargeCPU(program))
        trace = []
        for _ in range(20):
            step = runtime.step(); trace.append(step)
            if step.halt: break
        self.assertTrue(runtime.halted)
        self.assertTrue(any(step.pc == 40 for step in trace))
        self.assertEqual([step.acc for step in trace if step.output], [9])

    def test_future_event_is_not_visible(self):
        model = LargeCPU([('LOAD', 1), ('HALT', 0)])
        events = [encode_event(address=i, valid=True) for i in range(16)]
        events += [encode_event(position=16), encode_event(address=1, value=14, valid=True, position=17)]
        original = torch.tensor([events], dtype=torch.float32)
        changed = original.clone(); changed[0, -1, 17:21] = 0
        with torch.inference_mode():
            one, two = model(original, last_only=False), model(changed, last_only=False)
        self.assertTrue(torch.equal(one[:, :-1], two[:, :-1]))