"""Attention-backed nibble memory through a stock BeamDecoderLayer.

This probe uses constructed weights, a compact one-group Beam config and
host-encoded event vectors. Only memory reads are implemented in this decoder;
the layer's MLP is zeroed. It does not use pretrained Beam weights.
"""
import math
import torch
from transformers import BeamConfig
from transformers.cache_utils import DynamicCache
from transformers.models.beam.modeling_beam import BeamDecoderLayer, compute_attention_scaling
from .attention_cpu import bits


class DecoderMemory:
    def __init__(self, max_writes=128):
        if not 1 <= max_writes <= 1024: raise ValueError('max_writes must be 1..1024')
        self.maximum = max_writes
        self.writes = 0
        config = BeamConfig(hidden_size=128, intermediate_size=16, 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 = BeamDecoderLayer(config, 0).eval().requires_grad_(False)
        attn = self.layer.self_attn
        with torch.no_grad():
            for proj in (attn.q_proj, attn.k_proj, attn.v_proj, attn.o_proj, attn.gate_proj,
                         self.layer.mlp.gate_proj, self.layer.mlp.up_proj, self.layer.mlp.down_proj):
                proj.weight.zero_()
            for head in range(6):
                offset = 128*head
                for i in range(5): attn.q_proj.weight[offset+i, i] = 1
                attn.q_proj.weight[offset+5, 5] = 5
            for i in range(7): attn.k_proj.weight[i, 6+i] = 1
            for i in range(16):
                attn.v_proj.weight[i, 13+i] = 1
                attn.o_proj.weight[32+i, i] = 1
            attn.gate_bias.zero_()
            attn.q_norm.weight.fill_(16*max_writes*math.sqrt(30*6)/math.sqrt(128))
            attn.k_norm.weight.fill_(1)
        self.cache = DynamicCache(config=config)
        # Prefill zero-initialized RAM; earlier events cannot inspect future ones.
        rows = torch.stack([self._row(a, 0, 0, True) for a in range(16)]).unsqueeze(0)
        with torch.inference_mode(): self._forward(rows)

    def _row(self, address, value, position, valid):
        x = torch.zeros(128)
        x[:4] = torch.tensor([2*b-1 for b in bits(address, 4)])
        x[4:6] = 1
        x[6:10] = x[:4]
        p = position / self.maximum
        x[10], x[11], x[12] = p, 1 if valid else -1, math.sqrt(max(0, 1-p*p))
        x[13+value] = 1
        # The input layer's RMS normalization must not change V magnitudes.
        x[-1] = math.sqrt(128-float(x.square().sum()))
        return x

    def _forward(self, rows):
        old = self.cache.get_seq_length()
        length = rows.shape[1]
        mask = None
        if length > 1:
            mask = rows.new_zeros((1, 1, length, old+length))
            mask[0, 0, :, old:].masked_fill_(torch.ones(length, length, dtype=torch.bool).triu(1), float('-inf'))
        positions = torch.arange(old, old+length).unsqueeze(0)
        return self.layer(rows, position_embeddings=None, attention_mask=mask,
                          attention_scaling=compute_attention_scaling(positions),
                          past_key_values=self.cache, use_cache=True)

    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
        with torch.inference_mode(): self._forward(self._row(address, value, self.writes, True).reshape(1, 1, -1))

    def read(self, address):
        if not 0 <= address < 16: raise ValueError('expected nibble')
        with torch.inference_mode():
            output = self._forward(self._row(address, 0, self.writes, False).reshape(1, 1, -1))
        return int(output[0, 0, 32:48].argmax())