"""
Stage 2: The bigram language model — the simplest possible neural LM.

Originally this was ONE embedding table of shape (vocab_size, vocab_size):
row i a vector of raw next-token scores for "current token is i." Then a
second table added position information. Then a single self-attention Head.
Now the head is split into several narrower heads that run in parallel.

  - token_embedding_table:    one learned vector per vocab entry — the
                               token's identity, regardless of where it sits.
  - position_embedding_table: one learned vector per slot in the context
                               window — "this is position 3," regardless of
                               which token is there.
  - MultiHeadAttention:       n_head independent Heads, each working in a
                               narrower head_size = n_embd // n_head space,
                               run in parallel and concatenated back to
                               n_embd. Same total width as one big head, but
                               each head can specialize on a different kind
                               of relationship between positions instead of
                               averaging all of them into one.

    predict the next token; measure error with cross-entropy loss.

Everything a GPT does is this same objective with a smarter architecture.
Note the connection to the Word2vec slide from your lecture: like CBOW /
Skip-gram, this model learns embeddings. Position embeddings were a crude,
static fix for word order; self-attention is the real one — instead of a
fixed "position 3" vector, each token computes a data-dependent, weighted
average of what came before it.
"""

from __future__ import annotations

import torch
import torch.nn as nn
from torch.nn import functional as F


class Head(nn.Module):
    """One head of causal self-attention."""

    def __init__(self, n_embd: int, head_size: int, block_size: int) -> None:
        super().__init__()
        # Three separate linear projections of the same input x, each
        # (n_embd, head_size). No bias — these are pure projections, like
        # rotating/rescaling x into a new space, not adding an offset.
        self.key = nn.Linear(n_embd, head_size, bias=False)
        self.query = nn.Linear(n_embd, head_size, bias=False)
        self.value = nn.Linear(n_embd, head_size, bias=False)
        # A (block_size, block_size) lower-triangular matrix of 1s, e.g. for
        # block_size=4:
        #   1 0 0 0
        #   1 1 0 0
        #   1 1 1 0
        #   1 1 1 1
        # Row t has 1s in columns 0..t: "position t may attend to positions
        # 0..t." Registered as a buffer (not a parameter) so it moves with
        # .to(device) and is saved in state_dict, but is never trained.
        self.register_buffer("tril", torch.tril(torch.ones(block_size, block_size)))
        self.head_size = head_size

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        x: (B, T, n_embd) — the token+position embeddings for one batch.
        Returns (B, T, head_size): each position's attention-weighted
        summary of itself and everything before it.
        """
        B, T, _ = x.shape

        # Every position asks a question (query) and advertises an answer
        # (key). Both (B, T, n_embd) -> (B, T, head_size).
        q = self.query(x)
        k = self.key(x)

        # Compare every query to every key via dot product:
        # (B, T, head_size) @ (B, head_size, T) -> (B, T, T).
        # wei[b, i, j] = "how much position i's query matches position j's
        # key" — raw, unnormalized affinity, before masking or softmax.
        wei = q @ k.transpose(-2, -1)
        # Divide by sqrt(head_size) before softmax — see the explanation
        # printed by the demo script; in short, it keeps the variance of
        # these dot products at ~1 instead of growing with head_size, so
        # softmax doesn't saturate into a near one-hot distribution.
        wei = wei * self.head_size**-0.5

        # Causal mask: position i must not see position j > i (the future).
        # tril[:T, :T] == 0 marks the upper triangle (j > i); those entries
        # become -inf so softmax turns them into exactly 0.
        wei = wei.masked_fill(self.tril[:T, :T] == 0, float("-inf"))
        # Softmax over the last dim (T, the "which key" axis): every row i
        # becomes a probability distribution over positions 0..i that sums
        # to 1. Still (B, T, T).
        wei = F.softmax(wei, dim=-1)

        # Every position also advertises a value: (B, T, n_embd) ->
        # (B, T, head_size). This is what actually gets mixed together,
        # as opposed to key/query which only decide the mixing weights.
        v = self.value(x)
        # (B, T, T) @ (B, T, head_size) -> (B, T, head_size): position i's
        # output is the wei[i, :]-weighted average of every value vector at
        # positions 0..i.
        out = wei @ v
        return out


class MultiHeadAttention(nn.Module):
    """Several causal self-attention heads run in parallel, then combined."""

    def __init__(self, n_embd: int, n_head: int, block_size: int) -> None:
        super().__init__()
        # n_embd must split evenly across n_head so the concatenated output
        # lands back at exactly n_embd (e.g. 32 = 4 heads * 8).
        assert n_embd % n_head == 0, "n_embd must be divisible by n_head"
        head_size = n_embd // n_head
        self.heads = nn.ModuleList(
            [Head(n_embd, head_size, block_size) for _ in range(n_head)]
        )
        # Each head reads the same x (n_embd-wide) and produces its own
        # head_size-wide output. Concatenating n_head of those gets back to
        # n_embd; this projection lets the heads' outputs mix before moving
        # on, rather than just being stapled together.
        self.proj = nn.Linear(n_embd, n_embd)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Each h(x): (B, T, head_size). Concatenate along the last dim ->
        # (B, T, n_head * head_size) == (B, T, n_embd).
        out = torch.cat([h(x) for h in self.heads], dim=-1)
        return self.proj(out)


class FeedForward(nn.Module):
    """Per-position MLP: the "thinking" step after attention's "looking."

    Attention only moves information between positions via weighted
    averages of value vectors — it can't recombine features nonlinearly.
    This module runs independently on every position (no mixing across T)
    and gives the model a place to do that recombination. The inner
    dimension is widened 4x, as in the original Transformer paper, then
    projected back to n_embd so it can be added into the residual stream.
    """

    def __init__(self, n_embd: int) -> None:
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(n_embd, 4 * n_embd),
            nn.ReLU(),
            nn.Linear(4 * n_embd, n_embd),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x)


class Block(nn.Module):
    """One transformer block: communicate (attention), then compute (FFN).

    Pre-norm residual design: x = x + sublayer(LN(x)) for both sublayers.
    The "+" is what makes this residual — each sublayer adds a correction
    to x rather than replacing it, so during backprop gradients have a
    direct addition path (derivative 1) back through every block, no matter
    how deep the stack gets. Without it, gradients must flow through each
    block's full nonlinear transform in series, and repeated multiplication
    by small Jacobians is exactly the vanishing-gradient failure mode from
    the RNN slide. LayerNorm before each sublayer keeps that sublayer's
    input at a stable, zero-mean/unit-variance scale regardless of how
    large the residual stream has grown by this depth.
    """

    def __init__(self, n_embd: int, n_head: int, block_size: int) -> None:
        super().__init__()
        self.sa = MultiHeadAttention(n_embd, n_head, block_size)
        self.ffwd = FeedForward(n_embd)
        self.ln1 = nn.LayerNorm(n_embd)
        self.ln2 = nn.LayerNorm(n_embd)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x + self.sa(self.ln1(x))
        x = x + self.ffwd(self.ln2(x))
        return x


class BigramLanguageModel(nn.Module):
    def __init__(
        self,
        vocab_size: int,
        block_size: int,
        n_embd: int,
        n_head: int,
        n_layer: int,
    ) -> None:
        super().__init__()
        self.block_size = block_size
        # (vocab_size, n_embd): row i is token i's identity vector.
        self.token_embedding_table = nn.Embedding(vocab_size, n_embd)
        # (block_size, n_embd): row t is position t's vector. There are
        # only block_size rows because that's the longest context this
        # model will ever be asked about — see the crop in generate().
        self.position_embedding_table = nn.Embedding(block_size, n_embd)
        # n_layer stacked (attention, feedforward) blocks with residual
        # connections and layernorm — see Block's docstring for why those
        # two aren't decoration.
        self.blocks = nn.Sequential(
            *[Block(n_embd, n_head, block_size) for _ in range(n_layer)]
        )
        # One more LayerNorm after the last block. The residual stream's
        # scale is unconstrained as it accumulates additions across n_layer
        # blocks; this normalizes it once before the final projection reads
        # it, matching the scale lm_head was initialized to expect.
        self.ln_f = nn.LayerNorm(n_embd)
        # Project the final n_embd representation to vocab-sized logits.
        self.lm_head = nn.Linear(n_embd, vocab_size)

    def forward(
        self, idx: torch.Tensor, targets: torch.Tensor | None = None
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        """
        idx:     (B, T) batch of token-id sequences
        targets: (B, T) the same sequences shifted one position left,
                 i.e. "the correct next token at every position"

        Returns (logits, loss). Loss is None during generation.
        """
        B, T = idx.shape

        # (B, T) -> (B, T, n_embd): each token id looks up its identity vector.
        tok_emb = self.token_embedding_table(idx)
        # (T,) -> (T, n_embd): one vector per position, shared across the batch.
        pos_emb = self.position_embedding_table(torch.arange(T, device=idx.device))
        # (B, T, n_embd) + (T, n_embd): the position vectors broadcast over
        # the batch dim, so every sequence gets its own token vectors plus
        # the same T position vectors. Result: (B, T, n_embd).
        x = tok_emb + pos_emb
        # (B, T, n_embd) -> (B, T, n_embd): n_layer rounds of (attend, then
        # think), each wrapped in a residual connection.
        x = self.blocks(x)
        # Normalize the accumulated residual stream once before reading it.
        x = self.ln_f(x)
        # (B, T, n_embd) -> (B, T, vocab_size): project back to next-token scores.
        logits = self.lm_head(x)

        loss = None
        if targets is not None:
            B, T, C = logits.shape
            # cross_entropy wants (N, C) vs (N,), so flatten batch+time.
            loss = F.cross_entropy(logits.view(B * T, C), targets.view(B * T))
        return logits, loss

    @torch.no_grad()
    def generate(self, idx: torch.Tensor, max_new_tokens: int) -> torch.Tensor:
        """
        Autoregressive sampling: feed the sequence in, take the logits at
        the LAST position, turn them into probabilities, sample one token,
        append it, repeat. This loop is identical in shape to how a full
        GPT generates text — only the model in the middle changes.
        """
        for _ in range(max_new_tokens):
            # position_embedding_table has only block_size rows, so a
            # sequence longer than that would index out of range — crop to
            # the last block_size tokens before every forward pass.
            idx_cond = idx[:, -self.block_size :]
            logits, _ = self(idx_cond)         # (B, T, C)
            logits = logits[:, -1, :]          # last time step: (B, C)
            probs = F.softmax(logits, dim=-1)  # scores -> probabilities
            next_id = torch.multinomial(probs, num_samples=1)  # sample
            idx = torch.cat((idx, next_id), dim=1)
        return idx