"""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(64, d)                  # 3 in + 6 out max, generous
        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)

    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
        e = self.embed(seq) + 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)
        for i in range(T):
            for j in range(T):
                if j >= T_in and j > i:
                    blocked[i, j] = True
        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}