"""Sprout: a small GPT trained from scratch on TinyStories.

Modern-but-small transformer: RoPE, RMSNorm (pre-norm), SwiGLU MLP,
flash attention via F.scaled_dot_product_attention, untied heads.
"""
import math
import torch
import torch.nn as nn
import torch.nn.functional as F


def rope_cache(seq_len, head_dim, device, base=10000.0):
    """Precompute cos/sin tables for rotary embeddings [seq_len, head_dim/2]."""
    inv_freq = 1.0 / (base ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim))
    t = torch.arange(seq_len, device=device).float()
    freqs = torch.outer(t, inv_freq)                      # [seq, hd/2]
    return torch.cos(freqs), torch.sin(freqs)


def apply_rope(x, cos, sin):
    """x: [batch, heads, seq, head_dim]; cos/sin: [seq, head_dim/2]."""
    half = x.shape[-1] // 2
    x1, x2 = x[..., :half], x[..., half:]
    cos = cos[None, None, :, :].to(x.dtype)
    sin = sin[None, None, :, :].to(x.dtype)
    # rotate_half: swap halves with a sign flip
    out1 = x1 * cos - x2 * sin
    out2 = x2 * cos + x1 * sin
    return torch.cat([out1, out2], dim=-1)


class CausalSelfAttention(nn.Module):
    def __init__(self, dim, n_heads, dropout=0.0):
        super().__init__()
        assert dim % n_heads == 0
        self.n_heads = n_heads
        self.head_dim = dim // n_heads
        self.qkv = nn.Linear(dim, 3 * dim, bias=False)
        self.proj = nn.Linear(dim, dim, bias=False)
        self.dropout = dropout

    def forward(self, x, cos, sin):
        B, T, C = x.shape
        q, k, v = self.qkv(x).split(C, dim=2)
        shape = (B, T, self.n_heads, self.head_dim)
        q = q.view(shape).transpose(1, 2)  # [B, H, T, hd]
        k = k.view(shape).transpose(1, 2)
        v = v.view(shape).transpose(1, 2)
        q = apply_rope(q, cos, sin)
        k = apply_rope(k, cos, sin)
        y = F.scaled_dot_product_attention(
            q, k, v, is_causal=True,
            dropout_p=self.dropout if self.training else 0.0,
        )  # flash attention under the hood
        y = y.transpose(1, 2).contiguous().view(B, T, C)
        return self.proj(y)


class SwiGLU(nn.Module):
    """SwiGLU MLP; hidden dim set to ~8/3 of model dim rounded to multiple of 256."""

    def __init__(self, dim, dropout=0.0):
        super().__init__()
        hidden = max(256, int(dim * 8 / 3 / 256 + 0.5) * 256)
        self.w_gate = nn.Linear(dim, hidden, bias=False)
        self.w_up = nn.Linear(dim, hidden, bias=False)
        self.w_down = nn.Linear(hidden, dim, bias=False)
        self.drop = nn.Dropout(dropout)

    def forward(self, x):
        return self.drop(self.w_down(F.silu(self.w_gate(x)) * self.w_up(x)))


class Block(nn.Module):
    def __init__(self, dim, n_heads, dropout=0.0):
        super().__init__()
        self.n1 = nn.RMSNorm(dim, eps=1e-5)
        self.attn = CausalSelfAttention(dim, n_heads, dropout)
        self.n2 = nn.RMSNorm(dim, eps=1e-5)
        self.mlp = SwiGLU(dim, dropout)

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


class Sprout(nn.Module):
    def __init__(self, vocab_size, dim=512, n_layers=8, n_heads=8,
                 block_size=1024, dropout=0.0):
        super().__init__()
        self.block_size = block_size
        self.tok_emb = nn.Embedding(vocab_size, dim)
        self.drop = nn.Dropout(dropout)
        self.blocks = nn.ModuleList(Block(dim, n_heads, dropout) for _ in range(n_layers))
        self.norm_f = nn.RMSNorm(dim, eps=1e-5)
        self.head = nn.Linear(dim, vocab_size, bias=False)
        # init: normal(0.02), residual projections scaled by 1/sqrt(2*n_layers)
        self.apply(self._init_weights)
        for pn, p in self.named_parameters():
            if pn.endswith("proj.weight") or pn.endswith("w_down.weight"):
                nn.init.normal_(p, mean=0.0, std=0.02 / math.sqrt(2 * n_layers))

    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            nn.init.normal_(m.weight, mean=0.0, std=0.02)
            if m.bias is not None:
                nn.init.zeros_(m.bias)
        elif isinstance(m, nn.Embedding):
            nn.init.normal_(m.weight, mean=0.0, std=0.02)

    def forward(self, idx, targets=None):
        B, T = idx.shape
        assert T <= self.block_size, f"sequence {T} exceeds block size {self.block_size}"
        cos, sin = rope_cache(T, self.blocks[0].attn.head_dim, idx.device)
        x = self.drop(self.tok_emb(idx))
        for block in self.blocks:
            x = block(x, cos, sin)
        x = self.norm_f(x)
        if targets is not None:
            logits = self.head(x)
            return F.cross_entropy(
                logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1
            )
        # inference: only need the last position's logits
        return self.head(x[:, [-1], :])