"""
GPT-S2.5 — 5.11M-param subword language model (from scratch).
Custom GPT-2-style architecture: GQA + SwiGLU + RoPE + RMSNorm, weight-tied head.

This is NOT a transformers-native architecture. It is a self-contained
nn.Module you can load directly:

    import torch
    from model import Model
    m = Model()
    sd = torch.load("best.pt", map_location="cpu")["model"]   # or load safetensors
    m.load_state_dict(sd); m.eval()

    # generate (greedy):
    with torch.no_grad():
        x = torch.tensor([[1]])
        for _ in range(60):
            logits = m(x[:, -512:])[:, -1, :]
            x = torch.cat([x, logits.argmax(-1, keepdim=True)], dim=1)

Architecture (matches config.json exactly):
  vocab 8192 (BPE), n_embd 256, 4 layers, 8 q-heads / 2 kv-heads (GQA 4:1),
  head_dim 32, SwiGLU FFN intermediate 768, RoPE (base 10000), RMSNorm (eps 1e-6),
  pre-norm, no biases, weight-tied head (tok.weight reused as lm_head).
  Total: 5,114,112 parameters.
"""
import math
import torch
import torch.nn as nn
import torch.nn.functional as F

VOCAB      = 8192
N_EMBD      = 256
N_LAYERS    = 4
N_Q_HEADS   = 8
N_KV_HEADS  = 2
HEAD_DIM    = 32
INTER       = 768
MAX_SEQ     = 512
ROPE_BASE   = 10000.0

Q_DIM  = N_Q_HEADS * HEAD_DIM   # 256
KV_DIM = N_KV_HEADS * HEAD_DIM  # 64
assert Q_DIM == N_EMBD
assert N_Q_HEADS % N_KV_HEADS == 0


class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x):
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.weight


def precompute_rope(head_dim, max_seq, base):
    freqs = 1.0 / (base ** (torch.arange(0, head_dim, 2).float() / head_dim))
    t = torch.arange(max_seq, dtype=torch.float)
    ang = torch.outer(t, freqs)
    return ang.cos(), ang.sin()  # [T, head_dim/2]


def apply_rope(x, cos, sin):
    # x: [B, H, T, D]
    T = x.size(2)
    cos = cos[:T].view(1, 1, T, -1)
    sin = sin[:T].view(1, 1, T, -1)
    x1, x2 = x[..., 0::2], x[..., 1::2]
    out = torch.empty_like(x)
    out[..., 0::2] = x1 * cos - x2 * sin
    out[..., 1::2] = x2 * cos + x1 * sin
    return out


class Attention(nn.Module):
    def __init__(self):
        super().__init__()
        self.wq = nn.Linear(N_EMBD, Q_DIM, bias=False)
        self.wk = nn.Linear(N_EMBD, KV_DIM, bias=False)
        self.wv = nn.Linear(N_EMBD, KV_DIM, bias=False)
        self.wo = nn.Linear(Q_DIM, N_EMBD, bias=False)

    def forward(self, x, cos, sin):
        B, T, _ = x.shape
        q = self.wq(x).view(B, T, N_Q_HEADS, HEAD_DIM).transpose(1, 2)
        k = self.wk(x).view(B, T, N_KV_HEADS, HEAD_DIM).transpose(1, 2)
        v = self.wv(x).view(B, T, N_KV_HEADS, HEAD_DIM).transpose(1, 2)
        q = apply_rope(q, cos, sin)
        k = apply_rope(k, cos, sin)
        # GQA: repeat kv heads to match q heads
        rep = N_Q_HEADS // N_KV_HEADS
        k = k.repeat_interleave(rep, dim=1)
        v = v.repeat_interleave(rep, dim=1)
        y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        y = y.transpose(1, 2).contiguous().view(B, T, Q_DIM)
        return self.wo(y)


class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.wgate = nn.Linear(N_EMBD, INTER, bias=False)
        self.wup   = nn.Linear(N_EMBD, INTER, bias=False)
        self.wdown = nn.Linear(INTER, N_EMBD, bias=False)

    def forward(self, x):
        return self.wdown(F.silu(self.wgate(x)) * self.wup(x))


class Block(nn.Module):
    def __init__(self):
        super().__init__()
        self.ln1 = RMSNorm(N_EMBD)
        self.attn = Attention()
        self.ln2 = RMSNorm(N_EMBD)
        self.mlp = MLP()

    def forward(self, x, cos, sin):
        x = x + self.attn(self.ln1(x), cos, sin)
        x = x + self.mlp(self.ln2(x))
        return x


class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.tok = nn.Embedding(VOCAB, N_EMBD)
        self.blocks = nn.ModuleList([Block() for _ in range(N_LAYERS)])
        self.ln_f = RMSNorm(N_EMBD)
        self.cos, self.sin = precompute_rope(HEAD_DIM, MAX_SEQ, ROPE_BASE)

    def forward(self, idx, targets=None):
        B, T = idx.shape
        h = self.tok(idx)
        cos, sin = self.cos, self.sin
        for blk in self.blocks:
            h = blk(h, cos, sin)
        h = self.ln_f(h)
        logits = F.linear(h, self.tok.weight)  # weight-tied head
        if targets is not None:
            loss = F.cross_entropy(logits.view(-1, VOCAB), targets.view(-1))
            return loss
        return logits


if __name__ == "__main__":
    m = Model()
    total = sum(p.numel() for p in m.parameters())
    print(f"params = {total}  (expected 5114112)")
    assert total == 5114112, "param count mismatch"
    print("OK")