"""Two stock Beam decoder layers executing a compiled CPU transition.

Layer 0's SwiGLU MLP decodes the current PC into a memory address. Layer 1's
attention reads RAM, and its SwiGLU MLP computes the next CPU state. The host
serializes state/write events but does not look up ROM instructions during a
step. Weights and event format are constructed; this is not pretrained Beam.
"""
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, from_bits
from .hf_decoder_cpu import DecoderCPU
from .large_cpu import EventResult


class TwoLayerCPU(DecoderCPU):
    def __init__(self, program, max_writes=128):
        super().__init__(program, max_writes)
        config = BeamConfig(hidden_size=128, intermediate_size=256, num_hidden_layers=2,
                            num_attention_heads=6, num_key_value_heads=1, head_dim=128,
                            layer_types=['full_attention', 'full_attention'],
                            mlp_layer_types=['dense', 'dense'], attention_dropout=0.)
        config._attn_implementation = 'eager'
        self.first = BeamDecoderLayer(config, 0).eval().requires_grad_(False)
        self.layer.self_attn.layer_idx = 1
        self.layer.attn_residual_add_idx = 4
        self.layer.mlp_residual_add_idx = 5
        self.layer.self_attn.config = config
        # Layer 0 converts eight bipolar PC bits into four bipolar address bits.
        ids = torch.arange(256)
        patterns = (((ids[:, None] >> torch.arange(8)) & 1)*2-1).float()
        addresses = []
        for pc in range(256):
            op, arg = program[pc] if pc < len(program) else ('HALT', 0)
            addr = arg if op in {'LOAD', 'STORE', 'ADDM', 'SUBM'} else 0
            addresses.append([2*b-1 for b in bits(addr, 4)])
        with torch.no_grad():
            for proj in (self.first.self_attn.q_proj, self.first.self_attn.k_proj,
                         self.first.self_attn.v_proj, self.first.self_attn.o_proj,
                         self.first.self_attn.gate_proj, self.first.mlp.gate_proj,
                         self.first.mlp.up_proj, self.first.mlp.down_proj):
                proj.weight.zero_()
            self.first.self_attn.gate_bias.zero_()
            self.first.mlp.gate_proj.weight[:, 48:56] = 32 * patterns
            self.first.mlp.gate_proj.weight[:, 5] = -7.5 * 32
            self.first.mlp.up_proj.weight[:, 5] = 1
            self.first.mlp.down_proj.weight[:4, :] = torch.tensor(addresses, dtype=torch.float32).T
            # The first layer emits query-address bits with a different scale
            # from its unchanged fields; align them before layer-1 Q RMSNorm.
            for head in range(6):
                for i in range(4): self.layer.self_attn.q_proj.weight[head*128+i, i] *= 1/(4*math.sqrt(2))
            # Recalibrate the second layer MLP to its actual layer-1 residual
            # and normalization magnitudes: state=1/sqrt(3), read=sqrt(128/3).
            old_alpha, new_alpha = 1/math.sqrt(2), 1/math.sqrt(3)
            old_read, new_read = 8., math.sqrt(128/3)
            self.layer.mlp.gate_proj.weight[:, 48:60] *= old_alpha/new_alpha
            self.layer.mlp.gate_proj.weight[:, 5] *= old_alpha/new_alpha
            self.layer.mlp.gate_proj.weight[:, 32:48] *= old_read/new_read
            self.layer.mlp.up_proj.weight[:, 5] *= old_alpha/new_alpha
        self.cache = DynamicCache(config=config)
        self.pc = self.acc = self.writes = 0
        self.halted = False
        rows = torch.stack([self._event(a, 0, 0, True, 0, 0) for a in range(16)]).unsqueeze(0)
        with torch.inference_mode(): self._forward(rows)

    def _event(self, address, value, position, valid, pc, acc):
        # Query-address channels are zero. Layer 0 fills them from PC.
        x = torch.zeros(128)
        x[4:6] = 1
        x[6:10] = torch.tensor([2*b-1 for b in bits(address, 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
        x[48:60] = torch.tensor([2*b-1 for b in bits(pc, 8)+bits(acc, 4)])
        x[-1] = math.sqrt(128-float(x.square().sum()))
        return x

    def _forward(self, rows):
        if not hasattr(self, 'first'):
            return super()._forward(rows)
        old = self.cache.get_seq_length(layer_idx=1)
        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)
        kwargs = dict(position_embeddings=None, attention_mask=mask,
                      attention_scaling=compute_attention_scaling(positions),
                      past_key_values=self.cache, use_cache=True)
        return self.layer(self.first(rows, **kwargs), **kwargs)

    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._event(address, value, self.writes, True, self.pc, self.acc).reshape(1, 1, -1))

    def step(self):
        if self.halted: raise RuntimeError('machine is halted')
        old_acc = self.acc
        with torch.inference_mode():
            output = self._forward(self._event(0, 0, self.writes, False, self.pc, self.acc).reshape(1, 1, -1))[0, 0]
        if not torch.isfinite(output).all(): raise RuntimeError('non-finite output')
        address = from_bits((output[:4] > 0).int().tolist())
        decoded = output[60:90].reshape(15, 2).argmax(dim=-1).tolist()
        self.pc, self.acc = from_bits(decoded[:8]), from_bits(decoded[8:12])
        write, out, self.halted = map(bool, decoded[12:])
        if write: self.write(address, old_acc)
        return EventResult(self.pc, self.acc, write, out, self.halted, address, old_acc)