import math, torch, torch.nn as nn, torch.nn.functional as F

def sinusoidal_position_embedding(seq_len, dim, device):
    pe = torch.zeros(seq_len, dim, device=device)
    position = torch.arange(0, seq_len, device=device).unsqueeze(1)
    div_term = torch.exp(torch.arange(0, dim, 2, device=device) * (-math.log(10000.0) / dim))
    pe[:, 0::2] = torch.sin(position * div_term); pe[:, 1::2] = torch.cos(position * div_term)
    return pe

class CausalSelfAttention(nn.Module):
    def __init__(self, d_model, n_heads, dropout):
        super().__init__()
        assert d_model % n_heads == 0
        self.n_heads = n_heads; self.head_dim = d_model // n_heads
        self.qkv = nn.Linear(d_model, 3*d_model, bias=False)
        self.proj = nn.Linear(d_model, d_model, bias=False)
        self.attn_drop = nn.Dropout(dropout); self.resid_drop = nn.Dropout(dropout)
        self.register_buffer("mask", None, persistent=False)
    def forward(self, x):
        B, T, C = x.size()
        if self.mask is None or self.mask.size(0) < T:
            self.mask = torch.tril(torch.ones(T, T, device=x.device)).unsqueeze(0).unsqueeze(1)
        qkv = self.qkv(x); q, k, v = qkv.split(C, dim=2)
        def split(t): return t.view(B, T, self.n_heads, self.head_dim).transpose(1,2)
        q, k, v = map(split, (q,k,v))
        att = (q @ k.transpose(-2,-1)) / math.sqrt(self.head_dim)
        att = att.masked_fill(self.mask[:,:,:T,:T]==0, float('-inf'))
        att = F.softmax(att, dim=-1); att = self.attn_drop(att)
        y = att @ v; y = y.transpose(1,2).contiguous().view(B, T, C)
        return self.resid_drop(self.proj(y))

class MLP(nn.Module):
    def __init__(self, d_model, dropout):
        super().__init__()
        self.fc1 = nn.Linear(d_model, 4*d_model); self.fc2 = nn.Linear(4*d_model, d_model)
        self.drop = nn.Dropout(dropout)
    def forward(self, x):
        return self.drop(self.fc2(self.drop(F.gelu(self.fc1(x)))))

class Block(nn.Module):
    def __init__(self, d_model, n_heads, dropout):
        super().__init__()
        self.ln1 = nn.LayerNorm(d_model); self.attn = CausalSelfAttention(d_model, n_heads, dropout)
        self.ln2 = nn.LayerNorm(d_model); self.mlp  = MLP(d_model, dropout)
    def forward(self, x): x = x + self.attn(self.ln1(x)); return x + self.mlp(self.ln2(x))

class GPT(nn.Module):
    def __init__(self, vocab_size, d_model=512, n_heads=8, n_layers=8, n_positions=512, dropout=0.1):
        super().__init__()
        self.vocab_size = vocab_size; self.n_positions = n_positions
        self.tok_emb = nn.Embedding(vocab_size, d_model)
        self.pos_emb = nn.Parameter(torch.zeros(n_positions, d_model), requires_grad=False)
        self.drop = nn.Dropout(dropout)
        self.blocks = nn.ModuleList([Block(d_model, n_heads, dropout) for _ in range(n_layers)])
        self.ln_f = nn.LayerNorm(d_model); self.head = nn.Linear(d_model, vocab_size, bias=False)
        with torch.no_grad():
            self.pos_emb.copy_(sinusoidal_position_embedding(n_positions, d_model, device='cpu'))
    def forward(self, idx):
        B, T = idx.shape; tok = self.tok_emb(idx); pos = self.pos_emb[:T, :].unsqueeze(0).expand(B, -1, -1)
        x = self.drop(tok + pos)
        for blk in self.blocks: x = blk(x)
        x = self.ln_f(x); return self.head(x)