import random
import unittest
from beamboot.memory import AttentionMemory


class TestAttentionMemory(unittest.TestCase):
    def test_latest_write_and_missing(self):
        mem = AttentionMemory()
        self.assertIsNone(mem.read(3))
        mem.write(3, 12)
        mem.write(5, 99)
        mem.write(3, 231)
        self.assertEqual(mem.read_byte(3), 231)
        self.assertEqual(mem.read_byte(5), 99)
        self.assertIsNone(mem.read(2))

    def test_against_reference(self):
        randomizer = random.Random(42)
        mem, reference = AttentionMemory(temperature=0.00001), {}
        for _ in range(500):
            address = randomizer.randrange(16)
            if randomizer.random() < .55:
                value = randomizer.randrange(256)
                reference[address] = value
                mem.write(address, value)
            else:
                self.assertEqual(mem.read_byte(address), reference.get(address))

    def test_sliding_window_forgets(self):
        mem = AttentionMemory(window=2)
        mem.write(1, 42)
        mem.write(2, 10)
        mem.write(3, 20)
        self.assertIsNone(mem.read(1))
        self.assertEqual(mem.read_byte(2), 10)

    def test_softmax_is_not_exact(self):
        mem = AttentionMemory(temperature=1.0)
        mem.write(1, 0)
        mem.write(1, 255)
        self.assertNotEqual(mem.read_byte(1), 255)

    def test_validation(self):
        mem = AttentionMemory()
        with self.assertRaises(ValueError): mem.write(16, 1)
        with self.assertRaises(ValueError): mem.write(1, 256)
        with self.assertRaises(ValueError): mem.read(-1)


if __name__ == '__main__':
    unittest.main()