"""Binary-clock memory probe through Transformers' BeamAttention implementation.

This is a small constructed layer, not pretrained Beam. A host provides write
positions as bits. The eager Beam backend materializes attention scores in the
layer's dtype, unlike the separate SDPA probe's implementation.
"""
import math
import torch
from transformers import BeamConfig
from transformers.models.beam.modeling_beam import BeamAttention, compute_attention_scaling
from .attention_cpu import bits, from_bits


class HFBinaryMemory:
    def __init__(self, max_writes=128, dtype=torch.float32):
        if not 1 <= max_writes <= 2047: raise ValueError('max_writes must be 1..2047')
        self.maximum, self.dtype, self.writes = max_writes, dtype, 0
        self.width = max_writes.bit_length()
        self.clock = [2**i/(2*(max_writes+1)) for i in range(self.width)]
        config = BeamConfig(hidden_size=128, num_hidden_layers=1, num_attention_heads=6,
                            num_key_value_heads=1, head_dim=128, layer_types=['full_attention'],
                            mlp_layer_types=['dense'], attention_dropout=0.)
        config._attn_implementation = 'eager'
        self.layer = BeamAttention(config, 0).eval().to(dtype).requires_grad_(False)
        key_start = 6
        value_start = key_start + 5 + self.width
        with torch.no_grad():
            for m in (self.layer.q_proj, self.layer.k_proj, self.layer.v_proj,
                      self.layer.o_proj, self.layer.gate_proj): m.weight.zero_()
            for head in range(6):
                off = head*128
                for i in range(4): self.layer.q_proj.weight[off+i, i] = 1
                for i, coefficient in enumerate(self.clock): self.layer.q_proj.weight[off+4+i, 4] = coefficient
                self.layer.q_proj.weight[off+4+self.width, 4] = 2
            for i in range(5+self.width): self.layer.k_proj.weight[i, key_start+i] = 1
            for i in range(4):
                self.layer.v_proj.weight[i, value_start+i] = 1
                self.layer.o_proj.weight[i, i] = 2
            self.layer.gate_bias.zero_()
            q2 = 8 + sum(x*x for x in self.clock)
            k2 = 5 + self.width
            gain = 16*(max_writes+1)*math.sqrt(q2*k2)/math.sqrt(128)
            self.layer.q_norm.weight.fill_(gain)
            self.layer.k_norm.weight.fill_(1)
        self.rows = [self._row(a, 0, 0, True) for a in range(16)]

    def _row(self, address, value, position, valid):
        row = torch.zeros(128, dtype=self.dtype)
        row[:4] = torch.tensor([2*b-1 for b in bits(address, 4)], dtype=self.dtype)
        row[4] = 1
        key = [*(2*b-1 for b in bits(address, 4)),
               *(2*b-1 for b in bits(position, self.width)), 1 if valid else -1]
        row[6:6+len(key)] = torch.tensor(key, dtype=self.dtype)
        value_start = 6 + 5 + self.width
        row[value_start:value_start+4] = torch.tensor(bits(value, 4), dtype=self.dtype)
        return row

    def write(self, address, value):
        if not 0 <= address < 16 or not 0 <= value < 16: raise ValueError('expected nibbles')
        if self.writes >= self.maximum: raise ValueError('write budget exhausted')
        self.writes += 1
        self.rows.append(self._row(address, value, self.writes, True))

    def read(self, address):
        if not 0 <= address < 16: raise ValueError('expected nibble address')
        inputs = torch.stack(self.rows + [self._row(address, 0, self.writes, False)]).unsqueeze(0)
        length = inputs.shape[1]
        mask = torch.full((1, 1, length, length), float('-inf'), dtype=self.dtype).triu(1)
        positions = torch.arange(length).unsqueeze(0)
        with torch.inference_mode():
            out, _ = self.layer(inputs, position_embeddings=None, attention_mask=mask,
                                attention_scaling=compute_attention_scaling(positions))
        return from_bits((out[0, -1, :4] > .5).int().tolist())