import random
import unittest
import torch
from beamboot.tensor_memory import TensorMemory


class TensorMemoryTests(unittest.TestCase):
    def test_updates_match_reference(self):
        rng = random.Random(812)
        mem = TensorMemory(max_writes=256)
        reference = [0] * 16
        for _ in range(256):
            a, v = rng.randrange(16), rng.randrange(256)
            mem.write(a, v)
            reference[a] = v
            self.assertEqual(mem.read_byte(a), v)
        for a in range(16):
            self.assertEqual(mem.read_byte(a), reference[a])

    def test_window_loses_old_memory(self):
        mem = TensorMemory(window=2)
        mem.write(1, 42)
        mem.write(2, 20)
        mem.write(3, 30)
        self.assertNotEqual(mem.read_byte(1), 42)

    def test_write_limit_and_validation(self):
        mem = TensorMemory(max_writes=1)
        with self.assertRaises(ValueError): mem.write(16, 1)
        mem.write(0, 255)
        with self.assertRaises(ValueError): mem.write(0, 2)

    def test_bfloat16_is_measured_not_assumed(self):
        mem = TensorMemory(max_writes=4, dtype=torch.bfloat16)
        mem.write(1, 11)
        mem.write(1, 22)
        self.assertEqual(mem.read_byte(1), 22)