"""
A small character-level GPT (decoder-only transformer) in tinygrad.

Token + position embeddings -> a stack of pre-norm blocks (causal self-attention +
MLP, each added residually) -> a weight-tied head giving next-character logits.
"""
import math
from dataclasses import dataclass
from tinygrad import Tensor, nn


@dataclass
class GPTConfig:
    vocab_size: int = 162  # number of distinct characters (set from the data's meta.pkl)
    block_size: int = 256  # context window length (positions the model can see)
    n_layer: int = 6       # number of transformer blocks
    n_head: int = 6        # attention heads per block
    n_embd: int = 384      # width of the residual stream
    dropout: float = 0.2


class LayerNorm:
    """Layer norm with a learnable scale and no bias: normalize each position's
    feature vector to zero mean / unit variance, then rescale per channel."""
    def __init__(self, dim: int):
        self.weight = Tensor.ones(dim)

    def __call__(self, x: Tensor) -> Tensor:
        return x.layernorm() * self.weight


class CausalSelfAttention:
    """Multi-head self-attention where each position may only attend to earlier ones."""
    def __init__(self, cfg: GPTConfig):
        self.c_attn = nn.Linear(cfg.n_embd, 3 * cfg.n_embd, bias=False)  # q, k, v in one matmul
        self.c_proj = nn.Linear(cfg.n_embd, cfg.n_embd, bias=False)      # mix heads back together
        self.n_head, self.dropout = cfg.n_head, cfg.dropout

    def __call__(self, x: Tensor) -> Tensor:
        B, T, C = x.shape
        q, k, v = self.c_attn(x).split(C, dim=2)
        split_heads = lambda t: t.reshape(B, T, self.n_head, C // self.n_head).transpose(1, 2)
        q, k, v = split_heads(q), split_heads(k), split_heads(v)
        p = self.dropout if Tensor.training else 0.0
        y = q.scaled_dot_product_attention(k, v, is_causal=True, dropout_p=p)  # causal mask applied here
        y = y.transpose(1, 2).reshape(B, T, C)  # concatenate heads
        return self.c_proj(y).dropout(self.dropout)


class MLP:
    """Position-wise feed-forward network: expand 4x, GELU, project back."""
    def __init__(self, cfg: GPTConfig):
        self.c_fc = nn.Linear(cfg.n_embd, 4 * cfg.n_embd, bias=False)
        self.c_proj = nn.Linear(4 * cfg.n_embd, cfg.n_embd, bias=False)
        self.dropout = cfg.dropout

    def __call__(self, x: Tensor) -> Tensor:
        return self.c_proj(self.c_fc(x).gelu()).dropout(self.dropout)


class Block:
    """One transformer block: pre-norm attention and MLP, each added residually."""
    def __init__(self, cfg: GPTConfig):
        self.ln_1, self.attn = LayerNorm(cfg.n_embd), CausalSelfAttention(cfg)
        self.ln_2, self.mlp = LayerNorm(cfg.n_embd), MLP(cfg)

    def __call__(self, x: Tensor) -> Tensor:
        x = x + self.attn(self.ln_1(x))
        x = x + self.mlp(self.ln_2(x))
        return x


class GPT:
    """The full model: embeddings, a stack of blocks, and a tied output head."""
    def __init__(self, cfg: GPTConfig):
        self.cfg = cfg
        self.wte = nn.Embedding(cfg.vocab_size, cfg.n_embd)   # token embedding (reused as output head)
        self.wpe = nn.Embedding(cfg.block_size, cfg.n_embd)   # learned position embedding
        self.dropout = cfg.dropout
        self.blocks = [Block(cfg) for _ in range(cfg.n_layer)]
        self.ln_f = LayerNorm(cfg.n_embd)
        self._init_weights()

    def _init_weights(self):
        # Small-Gaussian init. Projections that feed a residual add are scaled down by
        # 1/sqrt(2*n_layer) so the variance of the residual stream stays stable with depth.
        randn = lambda shape, std: Tensor.randn(*shape) * std
        self.wte.weight = randn(self.wte.weight.shape, 0.02)
        self.wpe.weight = randn(self.wpe.weight.shape, 0.02)
        scaled = 0.02 / math.sqrt(2 * self.cfg.n_layer)
        for b in self.blocks:
            for lin in (b.attn.c_attn, b.mlp.c_fc):
                lin.weight = randn(lin.weight.shape, 0.02)
            for lin in (b.attn.c_proj, b.mlp.c_proj):
                lin.weight = randn(lin.weight.shape, scaled)

    def __call__(self, idx: Tensor) -> Tensor:
        B, T = idx.shape
        pos = Tensor.arange(T)
        x = (self.wte(idx) + self.wpe(pos)).dropout(self.dropout)  # combine token + position info
        for b in self.blocks:
            x = b(x)
        x = self.ln_f(x)
        return x.linear(self.wte.weight.transpose())  # tied output projection -> per-position logits

    def generate(self, idx: Tensor, max_new_tokens: int, temperature: float = 1.0,
                 top_k: int | None = None) -> Tensor:
        """Autoregressively extend `idx` (shape (B, T)) by sampling one token at a time.

        temperature scales the logits (lower = greedier); top_k, if set, restricts
        sampling to the k most likely next tokens. Returns the prompt plus new tokens.

        The context is always right-padded to block_size so every forward pass has the
        same shape (one kernel compile instead of one per length); causal masking means
        the padding never influences the real last position.
        """
        bs = self.cfg.block_size
        for _ in range(max_new_tokens):
            ctx = idx[:, -bs:]
            t = ctx.shape[1]
            if t < bs:
                ctx = ctx.cat(Tensor.zeros(ctx.shape[0], bs - t, dtype=idx.dtype), dim=1)
            logits = self(ctx)[:, t - 1, :] / temperature  # logits at the true last position
            if top_k:
                kth = logits.topk(min(top_k, logits.shape[-1]))[0][:, -1:]  # k-th largest per row
                logits = (logits < kth).where(-float("inf"), logits)        # mask the rest out
            next_id = logits.softmax(-1).multinomial(1).cast(idx.dtype)      # sample one token
            idx = idx.cat(next_id, dim=1).realize()  # run this step now (don't accumulate a lazy graph)
        return idx