"""Bounded attention memory with an exactly representable binary clock.

Like beam_probe, this is an isolated mathematical probe, not a Beam checkpoint.
Position is stored as bipolar bits so bfloat16 does not collapse adjacent values.
It still assumes an explicitly supplied, bounded write counter.
"""
import math
import torch
import torch.nn.functional as F
from .attention_cpu import bits


class BinaryTimeProbe:
    def __init__(self, max_writes=128, dtype=torch.float32):
        if not isinstance(max_writes, int) or not 1 <= max_writes <= 4095:
            raise ValueError('max_writes must be 1..4095')
        self.maximum, self.dtype = max_writes, dtype
        self.width = max_writes.bit_length()
        self.writes = 0
        self.keys, self.values = [], []
        # Each bit is exactly +/-1 in BF16. Normalized scalar clock is recovered
        # as a weighted sum, without quantizing the absolute position itself.
        self.weights = [2**i / (2 * (max_writes+1)) for i in range(self.width)]
        q2 = 4 + sum(w*w for w in self.weights)
        k2 = 4 + self.width
        self.gain = 16*(max_writes+1)*math.sqrt(q2*k2)/math.sqrt(128)
        for address in range(16): self._append(address, 0, 0)

    def _append(self, address, value, position):
        key = [*(2*b-1 for b in bits(address, 4)),
               *(2*b-1 for b in bits(position, self.width))]
        self.keys.append(key + [0.]*(128-len(key)))
        self.values.append(bits(value, 4) + [0.]*124)

    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 capacity exceeded')
        self.writes += 1
        self._append(address, value, self.writes)

    @staticmethod
    def norm(x):
        return x * torch.rsqrt(x.float().square().mean(dim=-1, keepdim=True) + 1e-6).to(x.dtype)

    def read(self, address):
        if not 0 <= address < 16: raise ValueError('expected nibble address')
        q = [*(2*b-1 for b in bits(address, 4)), *self.weights]
        query = torch.tensor(q + [0.]*(128-len(q)), dtype=self.dtype)
        keys = torch.tensor(self.keys, dtype=self.dtype)
        values = torch.tensor(self.values, dtype=self.dtype)
        query = self.norm(query) * self.gain
        keys = self.norm(keys)
        position = len(self.keys)
        query = query * ((1 + (position+1)/2048)**.1)
        q = query.reshape(1, 1, 1, 128).expand(1, 6, 1, 128)
        k = keys.reshape(1, 1, -1, 128)
        v = values.reshape(1, 1, -1, 128)
        out = F.scaled_dot_product_attention(q, k, v, enable_gqa=True)[0, 0, 0, :4]
        return sum((float(x) > .5) << i for i, x in enumerate(out))