"""
model.py — Path B architecture for the ~250M SLM base.

Modifications from build-nanogpt's GPT:
  1. GQA (grouped-query attention) — fewer KV heads than query heads => small KV cache
  2. RoPE (rotary position embeddings) — replaces learned wpe; extendable context
  3. Tied embeddings — kept from baseline
  4. Vocab size is a config field (set it to your trained tokenizer's size, e.g. 32768)

Kept simple (LayerNorm + GELU) on purpose for a low-risk first real model.
This file only DEFINES the model — it never trains on import, so it's safe to
`from model import GPT, GPTConfig` from any script.
"""
import math
from dataclasses import dataclass
import torch
import torch.nn as nn
from torch.nn import functional as F

# -----------------------------------------------------------------------------
# RoPE helpers

def build_rope_cache(seq_len, head_dim, device, base=10000.0):
    """Precompute cos/sin tables for rotary embeddings. Shape: (seq_len, head_dim)."""
    assert head_dim % 2 == 0, "head_dim must be even for RoPE"
    theta = 1.0 / (base ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim))
    positions = torch.arange(seq_len, device=device).float()
    freqs = torch.outer(positions, theta)                 # (seq_len, head_dim/2)
    emb = torch.cat([freqs, freqs], dim=-1)               # (seq_len, head_dim)
    return emb.cos(), emb.sin()

def rotate_half(x):
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat((-x2, x1), dim=-1)

def apply_rope(q, k, cos, sin):
    # q, k: (B, n_head, T, head_dim); cos, sin: (T, head_dim)
    cos = cos.unsqueeze(0).unsqueeze(0)   # (1, 1, T, head_dim)
    sin = sin.unsqueeze(0).unsqueeze(0)
    q_rot = (q * cos) + (rotate_half(q) * sin)
    k_rot = (k * cos) + (rotate_half(k) * sin)
    return q_rot, k_rot

# -----------------------------------------------------------------------------
# Grouped-Query Attention with RoPE

class CausalSelfAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        assert config.n_embd % config.n_head == 0
        assert config.n_head % config.n_kv_head == 0, "n_head must be divisible by n_kv_head"
        self.n_head = config.n_head
        self.n_kv_head = config.n_kv_head
        self.n_embd = config.n_embd
        self.head_dim = config.n_embd // config.n_head
        self.n_rep = self.n_head // self.n_kv_head   # how many query heads share each KV head

        # Q projects to full n_head; K and V project to only n_kv_head => smaller KV
        self.q_proj = nn.Linear(config.n_embd, self.n_head * self.head_dim, bias=False)
        self.k_proj = nn.Linear(config.n_embd, self.n_kv_head * self.head_dim, bias=False)
        self.v_proj = nn.Linear(config.n_embd, self.n_kv_head * self.head_dim, bias=False)
        self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=False)
        self.c_proj.NANOGPT_SCALE_INIT = 1

    def forward(self, x, cos, sin):
        B, T, C = x.size()
        q = self.q_proj(x).view(B, T, self.n_head, self.head_dim).transpose(1, 2)     # (B, nh, T, hd)
        k = self.k_proj(x).view(B, T, self.n_kv_head, self.head_dim).transpose(1, 2)  # (B, nkv, T, hd)
        v = self.v_proj(x).view(B, T, self.n_kv_head, self.head_dim).transpose(1, 2)  # (B, nkv, T, hd)

        # apply rotary embeddings to q and k
        q, k = apply_rope(q, k, cos[:T], sin[:T])

        # expand KV heads to match query heads (GQA): repeat each KV head n_rep times
        k = k.repeat_interleave(self.n_rep, dim=1)   # (B, nh, T, hd)
        v = v.repeat_interleave(self.n_rep, dim=1)   # (B, nh, T, hd)

        y = F.scaled_dot_product_attention(q, k, v, is_causal=True)  # flash attention
        y = y.transpose(1, 2).contiguous().view(B, T, C)
        return self.c_proj(y)

class MLP(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd, bias=False)
        self.gelu = nn.GELU(approximate='tanh')
        self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd, bias=False)
        self.c_proj.NANOGPT_SCALE_INIT = 1

    def forward(self, x):
        return self.c_proj(self.gelu(self.c_fc(x)))

class Block(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.ln_1 = nn.LayerNorm(config.n_embd)
        self.attn = CausalSelfAttention(config)
        self.ln_2 = nn.LayerNorm(config.n_embd)
        self.mlp = MLP(config)

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

# -----------------------------------------------------------------------------

@dataclass
class GPTConfig:
    block_size: int = 2048      # context length
    vocab_size: int = 32768     # SET to your trained tokenizer's size (incl. FIM/special tokens)
    n_layer: int = 24
    n_head: int = 16            # query heads
    n_kv_head: int = 4          # KV heads (GQA); n_head/n_kv_head = 4 query heads per KV head
    n_embd: int = 1024

class GPT(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.transformer = nn.ModuleDict(dict(
            wte=nn.Embedding(config.vocab_size, config.n_embd),
            h=nn.ModuleList([Block(config) for _ in range(config.n_layer)]),
            ln_f=nn.LayerNorm(config.n_embd),
        ))
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
        # tied embeddings
        self.transformer.wte.weight = self.lm_head.weight

        # RoPE cache (built lazily on first forward, cached on the module)
        self.register_buffer("rope_cos", None, persistent=False)
        self.register_buffer("rope_sin", None, persistent=False)

        self.apply(self._init_weights)

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            std = 0.02
            if hasattr(module, 'NANOGPT_SCALE_INIT'):
                std *= (2 * self.config.n_layer) ** -0.5
            torch.nn.init.normal_(module.weight, mean=0.0, std=std)
            if module.bias is not None:
                torch.nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def _ensure_rope(self, T, device):
        head_dim = self.config.n_embd // self.config.n_head
        if self.rope_cos is None or self.rope_cos.size(0) < T or self.rope_cos.device != device:
            cos, sin = build_rope_cache(max(T, self.config.block_size), head_dim, device)
            self.rope_cos, self.rope_sin = cos, sin

    def forward(self, idx, targets=None):
        B, T = idx.size()
        assert T <= self.config.block_size, f"sequence length {T} > block_size {self.config.block_size}"
        self._ensure_rope(T, idx.device)
        cos, sin = self.rope_cos.to(idx.device), self.rope_sin.to(idx.device)

        x = self.transformer.wte(idx)     # (B, T, n_embd) — no positional embedding added; RoPE handles it
        for block in self.transformer.h:
            x = block(x, cos, sin)
        x = self.transformer.ln_f(x)
        logits = self.lm_head(x)
        loss = None
        if targets is not None:
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
        return logits, loss

    def num_params(self):
        n = sum(p.numel() for p in self.parameters())
        # subtract tied lm_head (shares wte weight) to avoid double counting
        n -= self.lm_head.weight.numel()
        return n

# -----------------------------------------------------------------------------
# quick self-test: prints param count so you can tune the config to ~250M
if __name__ == "__main__":
    cfg = GPTConfig()
    model = GPT(cfg)
    print(f"config: n_layer={cfg.n_layer}, n_embd={cfg.n_embd}, "
          f"n_head={cfg.n_head}, n_kv_head={cfg.n_kv_head}, vocab={cfg.vocab_size}")
    print(f"total parameters: {model.num_params()/1e6:.1f}M")
    # tiny forward sanity check on CPU
    x = torch.randint(0, cfg.vocab_size, (2, 128))
    logits, _ = model(x)
    print(f"forward OK — logits shape {tuple(logits.shape)}")