import random
import unittest
import torch
from beamboot.binary_time_probe import BinaryTimeProbe
from beamboot.hf_binary_memory import HFBinaryMemory
from beamboot.causal_lm import CausalMachine


class BinaryPrecisionTests(unittest.TestCase):
    def test_bf16_binary_clock_and_retention(self):
        memory = BinaryTimeProbe(512, torch.bfloat16)
        rng = random.Random(417)
        reference = [0]*16
        for i in range(512):
            address, value = rng.randrange(16), rng.randrange(16)
            memory.write(address, value)
            reference[address] = value
            self.assertEqual(memory.read(address), value, f'write {i+1}')
        for address, value in enumerate(reference):
            self.assertEqual(memory.read(address), value)

    def test_hf_beam_attention_sdpa_bf16(self):
        memory = HFBinaryMemory(64, torch.bfloat16)
        memory.layer.config._attn_implementation = 'sdpa'
        rng = random.Random(417)
        for _ in range(64):
            address, value = rng.randrange(16), rng.randrange(16)
            memory.write(address, value)
            self.assertEqual(memory.read(address), value)

    def test_known_unsafe_runtime_is_rejected(self):
        with self.assertRaisesRegex(ValueError, 'known incorrect'):
            CausalMachine._validate_precision(torch.bfloat16, 'eager')
        CausalMachine._validate_precision(torch.bfloat16, 'sdpa')