"""Experimental CPU transition inside a stock BeamDecoderLayer.

Both RAM retrieval and ROM-specific transitions execute in stock Beam decoder
attention and SwiGLU MLP weights. The host still selects the ROM operand/address
from the current PC and serializes discrete state and write events. This is not
an autoregressive LM, a trained checkpoint, or an embedded 501B Beam model.
"""
import math
import torch
from transformers.models.beam.modeling_beam import BeamMLP
from .attention_cpu import bits, from_bits
from .hf_decoder_memory import DecoderMemory
from .large_cpu import validate, transition, EventResult


class DecoderCPU(DecoderMemory):
    def __init__(self, program, max_writes=128):
        validate(program)
        self.program = list(program)
        super().__init__(max_writes)
        # Each neuron recognizes one of 256*16*16 (PC, ACC, RAM) states.
        # This is a large, exponentially scaling truth table, not a practical CPU.
        n = 65536
        mlp = BeamMLP(self.layer.self_attn.config, intermediate_size=n)
        ids = torch.arange(n, dtype=torch.int64)
        pc = ids & 255
        acc = (ids >> 8) & 15
        mem = (ids >> 12) & 15
        state = (((ids[:, None] >> torch.arange(12)) & 1)*2-1).float()
        self.layer.mlp = mlp.eval().requires_grad_(False)
        # After the attention residual, a bipolar input bit has magnitude 1/sqrt(2)
        # and the one-hot memory response has magnitude 8. Layer norms preserve
        # these ratios because the two subspaces are orthogonal and fixed-norm.
        alpha = 1 / math.sqrt(2)
        temperature = 32.0
        transitions = []
        for i in range(n):
            npc, nacc, write, output, halt = transition(program, int(pc[i]), int(acc[i]), int(mem[i]))
            transitions.append(bits(npc, 8) + bits(nacc, 4) + [write, output, halt])
        transitions = torch.tensor(transitions, dtype=torch.int64)
        with torch.no_grad():
            mlp.gate_proj.weight.zero_(); mlp.up_proj.weight.zero_(); mlp.down_proj.weight.zero_()
            mlp.gate_proj.weight[:, 48:60] = state * (temperature / alpha)
            mlp.gate_proj.weight[ids, 32 + mem] = temperature / 8
            mlp.gate_proj.weight[:, 5] = -12.5 * temperature / alpha
            mlp.up_proj.weight[:, 5] = 1 / alpha
            # Two output channels per bit; exactly one is active per transition.
            for i in range(15):
                mlp.down_proj.weight[60 + 2*i, ids] = (transitions[:, i] == 0).float()
                mlp.down_proj.weight[61 + 2*i, ids] = (transitions[:, i] == 1).float()
        self.pc = 0
        self.acc = 0
        self.halted = False

    def _cpu_row(self, address, pc, acc):
        x = self._row(address, 0, self.writes, False)
        state = bits(pc, 8) + bits(acc, 4)
        x[48:60] = torch.tensor([2*b-1 for b in state], dtype=torch.float32)
        x[-1] = 0
        x[-1] = math.sqrt(128-float(x.square().sum()))
        return x

    def device_write(self, address, value):
        self.write(address, value)

    def step(self):
        if self.halted: raise RuntimeError('machine is halted')
        # The address decoder still runs on the host: the model's transition
        # is compiled, but this protocol is not yet a self-contained decoder.
        op, operand = self.program[self.pc] if self.pc < len(self.program) else ('HALT', 0)
        address = operand if op in {'LOAD', 'STORE', 'ADDM', 'SUBM'} else 0
        old_acc = self.acc
        with torch.inference_mode():
            output = self._forward(self._cpu_row(address, self.pc, self.acc).reshape(1, 1, -1))[0, 0]
        pairs = output[60:90].reshape(15, 2)
        if not torch.isfinite(pairs).all(): raise RuntimeError('non-finite output')
        decoded = pairs.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)