"""Bounded cache memory using PyTorch's scaled-dot-product attention.

This is a constructed attention experiment, not Beam's attention implementation.
"""
import torch
import torch.nn.functional as F


class TensorMemory:
    def __init__(self, address_bits=4, max_writes=256, window=None, dtype=torch.float32):
        if address_bits < 1 or max_writes < 1 or (window is not None and window < 1):
            raise ValueError('invalid configuration')
        self.address_bits = address_bits
        self.max_writes = max_writes
        self.window = window
        self.dtype = dtype
        self.keys = []
        self.values = []
        # Initial zeros establish all addresses, removing missing-key checks.
        for address in range(2**address_bits):
            self._append(address, 0, 0)
        self.writes = 0

    def _bits(self, address):
        if not isinstance(address, int) or not 0 <= address < 2**self.address_bits:
            raise ValueError('address out of range')
        return [1.0 if address & (1 << i) else -1.0 for i in range(self.address_bits)]

    def _append(self, address, value, position):
        # Matching an address dominates any possible recency difference, while
        # a recency difference of 16 makes the latest match dominate softmax.
        self.keys.append(self._bits(address) + [float(position)])
        self.values.append([float(value)])

    def write(self, address, value):
        if self.writes >= self.max_writes:
            raise ValueError('write budget exhausted')
        if not isinstance(value, int) or not 0 <= value <= 255:
            raise ValueError('value must be a byte')
        self._bits(address)
        self.writes += 1
        self._append(address, value, self.writes)

    def read(self, address):
        bits = self._bits(address)
        # An address mismatch loses at least 2 * (16 * (max_writes + 1)),
        # exceeding the full position score span. PyTorch scale=1 makes these
        # logits explicit, without an additive score mask.
        query = bits + [1.0]
        keys = self.keys[-self.window:] if self.window else self.keys
        values = self.values[-self.window:] if self.window else self.values
        q = torch.tensor(query, dtype=self.dtype).reshape(1, 1, 1, -1)
        k = torch.tensor(keys, dtype=self.dtype).reshape(1, 1, len(keys), -1)
        v = torch.tensor(values, dtype=self.dtype).reshape(1, 1, len(values), 1)
        q[..., :-1] *= 16 * (self.max_writes + 1)
        q[..., -1] *= 16
        with torch.no_grad():
            return float(F.scaled_dot_product_attention(q, k, v, scale=1.0).item())

    def read_byte(self, address):
        return round(self.read(address))