"""Transformer baseline: GPT-style causal decoder over [input digits | output digits].

Fixed d_model (=128) matching the RNN arm. Parameter counts are logged per run but
NOT gated to parity (design/preregistration.md: weight sharing is the studied variable).
"""
import torch
import torch.nn as nn

from src.config import Config
from src.model_api import PrimeModel


class CausalBlock(nn.Module):
    def __init__(self, d: int, heads: int):
        super().__init__()
        self.ln1 = nn.LayerNorm(d)
        self.ln2 = nn.LayerNorm(d)
        self.attn = nn.MultiheadAttention(d, heads, batch_first=True)
        self.mlp = nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(), nn.Linear(4 * d, d))

    def forward(self, x: torch.Tensor, attn_mask: torch.Tensor, key_pad: torch.Tensor) -> torch.Tensor:
        a, _ = self.attn(self.ln1(x), self.ln1(x), self.ln1(x),
                         attn_mask=attn_mask, key_padding_mask=key_pad, need_weights=False)
        x = x + a
        x = x + self.mlp(self.ln2(x))
        return x


class TransformerBaseline(PrimeModel):
    def __init__(self, cfg: Config):
        super().__init__()
        self.cfg = cfg
        d = cfg.d_model
        self.embed = nn.Embedding(cfg.vocab + 1, d)     # +1 row = pad
        self.pos = nn.Embedding(256, d)                 # supports sequences up to 256
        self.blocks = nn.ModuleList([CausalBlock(d, cfg.n_heads) for _ in range(cfg.n_layers)])
        self.ln_f = nn.LayerNorm(d)
        self.head = nn.Linear(d, cfg.vocab)
        self.register_buffer(
            "causal_triu",
            torch.triu(torch.ones(256, 256, dtype=torch.bool), diagonal=1),
            persistent=False,
        )

    def forward(self, x: torch.Tensor, y_in: torch.Tensor) -> dict:
        cfg = self.cfg
        B, T_in = x.shape
        T_out = y_in.shape[1]
        seq = torch.cat([x, y_in], dim=1)               # (B, T_in+T_out)
        T = seq.shape[1]
        device = seq.device
        
        # Token embeddings
        tok_embed = self.embed(seq)
        if cfg.scratch_mode == "random_frozen":
            # detach embeddings for tokens 'a'..'p' (IDs 13..28)
            is_rand = (seq >= 13) & (seq <= 28)
            tok_embed = torch.where(is_rand.unsqueeze(-1), tok_embed.detach(), tok_embed)
        elif cfg.scratch_mode == "random_noise":
            # inject fresh Gaussian noise at NOISE_SLOT_ID (ID 33)
            is_noise = (seq == 33)
            noise = torch.randn_like(tok_embed)
            tok_embed = torch.where(is_noise.unsqueeze(-1), noise, tok_embed)

        e = tok_embed + self.pos(torch.arange(T, device=device)).unsqueeze(0)
        # causal mask: input positions (j < T_in) fully visible; output positions causal
        # (True = blocked, per torch.nn.MultiheadAttention bool convention)
        blocked = torch.zeros(T, T, dtype=torch.bool, device=device)
        blocked[:, T_in:] = self.causal_triu[:T, T_in:T]
        key_pad = seq == cfg.pad_id                       # (B,T) True = ignore
        h = e
        for blk in self.blocks:
            h = blk(h, blocked, key_pad)
        # logits at position p predict token p+1 -> positions [T_in-1, T_in+T_out-2] predict y
        logits = self.head(self.ln_f(h))[:, T_in - 1: T_in - 1 + T_out]
        return {"logits": logits, "halt_steps": None}