"""Causal attention CPU over a stream of typed event vectors.

This standalone constructed decoder recomputes projected keys/values from its
input event stream. Its attention and MLP weights execute the CPU. The host
serializes model-emitted states and writes into event vectors, runs I/O and
stops at halt. It is not a Beam checkpoint or Hugging Face Transformers model.
"""
import torch
from torch import nn
import torch.nn.functional as F
from .attention_cpu import AttentionCPU, STATE_BITS, Step, bits, from_bits

# Input event vector: PC[5], ACC[4], write-valid, address[4], data[4], position.
FEATURES = 19
VALID, ADDRESS, DATA, POSITION = 9, 10, 14, 18


def event(state=None, address=0, value=0, valid=False, position=0):
    state = [0] * STATE_BITS if state is None else list(state)
    if len(state) != STATE_BITS or any(x not in (0, 1) for x in state):
        raise ValueError('state must contain nine binary values')
    if not 0 <= address < 16 or not 0 <= value < 16 or position < 0:
        raise ValueError('invalid event')
    return state + [int(valid)] + bits(address, 4) + bits(value, 4) + [position]


class SequenceCPU(nn.Module):
    def __init__(self, program, max_events=1024):
        super().__init__()
        if not 16 < max_events <= 8192: raise ValueError('max_events must be 17..8192')
        self.max_events = max_events
        core = AttentionCPU(program)
        self.pc_detect, self.address = core.pc_detect, core.address
        self.detect, self.result = core.detect, core.result
        self.k_proj = nn.Linear(FEATURES, 6)
        self.v_proj = nn.Linear(FEATURES, 4, bias=False)
        with torch.no_grad():
            self.k_proj.weight.zero_(); self.k_proj.bias.zero_()
            for i in range(4):
                self.k_proj.weight[i, ADDRESS+i] = 2
                self.k_proj.bias[i] = -1
            self.k_proj.weight[4, POSITION] = 1
            self.k_proj.weight[5, VALID] = 2
            self.k_proj.bias[5] = -1
            self.v_proj.weight.zero_()
            for i in range(4): self.v_proj.weight[i, DATA+i] = 1
        self.requires_grad_(False)

    def forward(self, inputs, last_only=False):
        """Input shape [batch,time,19]; output shape [batch,time,20].

        With last_only, output time dimension is one and all supplied events
        are in its past. Full mode uses an explicit causal attention mask.
        """
        if inputs.ndim != 3 or inputs.shape[-1] != FEATURES: raise ValueError('invalid input shape')
        if not 0 < inputs.shape[1] <= self.max_events: raise ValueError('event budget exceeded')
        selected = inputs[:, -1:, :] if last_only else inputs
        state = selected[..., :STATE_BITS]
        address = self.address(torch.relu(self.pc_detect(state[..., :5])))
        base = 16 * (self.max_events + 1)
        query = torch.cat(((2 * address - 1) * base,
                           torch.ones_like(address[..., :1]) * 16,
                           torch.ones_like(address[..., :1]) * base * 5), dim=-1)
        keys, values = self.k_proj(inputs), self.v_proj(inputs)
        memory = F.scaled_dot_product_attention(query.unsqueeze(1), keys.unsqueeze(1), values.unsqueeze(1),
                                                scale=1., is_causal=not last_only).squeeze(1)
        read_bits = 2 * (torch.relu(memory - .25) - torch.relu(memory - .75))
        result = self.result(torch.relu(self.detect(torch.cat((state, read_bits), dim=-1))))
        return torch.cat((result, address, state[..., 5:]), dim=-1)


class SequenceRuntime:
    """Serializes event vectors and feeds back predicted machine state."""
    def __init__(self, model):
        self.model = model
        self.state = torch.zeros(STATE_BITS)
        self.events = [event(address=i, valid=True) for i in range(16)]
        self.halted = False

    def _append(self, item):
        if len(self.events) >= self.model.max_events: raise RuntimeError('event budget exhausted')
        self.events.append(item)

    def device_write(self, address, value):
        if not isinstance(address, int) or not isinstance(value, int) or not 0 <= address < 16 or not 0 <= value < 16:
            raise ValueError('device address and value must be 0..15')
        self._append(event(address=address, value=value, valid=True, position=len(self.events)))

    def step(self):
        if self.halted: raise RuntimeError('machine is halted')
        self._append(event(self.state.tolist(), position=len(self.events)))
        with torch.inference_mode():
            output = self.model(torch.tensor([self.events], dtype=torch.float32), last_only=True)[0, 0]
        if not torch.isfinite(output).all(): raise RuntimeError('non-finite circuit output')
        emitted = (output > .5).to(torch.int64).tolist()
        new_state = emitted[:STATE_BITS]
        write, out, halt = map(bool, emitted[STATE_BITS:STATE_BITS+3])
        address = from_bits(emitted[STATE_BITS+3:STATE_BITS+7])
        value = from_bits(emitted[STATE_BITS+7:STATE_BITS+11])
        if write: self.device_write(address, value)
        self.state = torch.tensor(new_state, dtype=torch.float32)
        self.halted = halt
        return Step(from_bits(new_state[:5]), from_bits(new_state[5:]), write, out, halt, address, value)