"""Generates 01_build_gpt.ipynb. Run once with any Python (uses only stdlib json)."""
import json
import os

HERE = os.path.dirname(os.path.abspath(__file__))

cells = []

def md(src):
    cells.append({"cell_type": "markdown", "metadata": {}, "source": src})

def code(src):
    cells.append({"cell_type": "code", "metadata": {}, "execution_count": None,
                  "outputs": [], "source": src})

md("""# Let's build GPT — from scratch

We are rebuilding a tiny GPT on the *tiny Shakespeare* dataset, following Karpathy's
"Let's build GPT" but with extra explanation at each step.

**Before running:** in VS Code, click the kernel picker (top-right) and choose the
interpreter at `gpt_from_scratch/.venv`. That venv has CUDA PyTorch installed.

The roadmap inside this notebook:
1. Data + tokenization
2. Batching (the input/target shapes)
3. A Bigram baseline (simplest possible language model) + training loop
4. *(next session)* Self-attention — the heart of the Transformer
""")

md("""## 0. Setup check

First confirm PyTorch can actually see your GPU. If `CUDA available` is False,
stop — we're accidentally on the CPU build again.""")

code("""import torch
import torch.nn as nn
from torch.nn import functional as F

print("torch:", torch.__version__)
print("CUDA available:", torch.cuda.is_available())
device = "cuda" if torch.cuda.is_available() else "cpu"
print("device:", device)
if device == "cuda":
    print("GPU:", torch.cuda.get_device_name(0))""")

md("""## 1. The data

A language model learns to predict the *next* piece of text given what came before.
So all we need is a big pile of text. Ours is ~1MB of Shakespeare.""")

code("""with open("input.txt", "r", encoding="utf-8") as f:
    text = f.read()

print("length in characters:", len(text))
print("----")
print(text[:250])""")

md("""## 2. Tokenization

A neural net only eats **numbers**, not letters. *Tokenization* = deciding how to
chop text into pieces ("tokens") and mapping each piece to an integer.

Real GPTs use *sub-word* tokens (e.g. "ing", "tion") via an algorithm called BPE.
We'll start with the simplest possible scheme: **one token = one character**.

- Pro: trivially simple, tiny vocabulary.
- Con: sequences get long (every letter is a step), so the model has to work harder.

> Analogy: tokenization is choosing your alphabet. Characters = small alphabet, long
> words. Sub-words = bigger alphabet, shorter words. GPT picks the bigger alphabet.""")

code("""# all the unique characters that occur in the text = our vocabulary
chars = sorted(list(set(text)))
vocab_size = len(chars)
print("".join(chars))
print("vocab_size:", vocab_size)

# build the two lookup tables: char->int (stoi) and int->char (itos)
stoi = {ch: i for i, ch in enumerate(chars)}
itos = {i: ch for i, ch in enumerate(chars)}
encode = lambda s: [stoi[c] for c in s]          # string -> list of ints
decode = lambda l: "".join(itos[i] for i in l)   # list of ints -> string

print(encode("hello there"))
print(decode(encode("hello there")))""")

md("""## 3. Encode the whole dataset into a tensor

We turn the entire text into one long 1-D tensor of integers, then hold out the last
10% as a validation set (to check the model isn't just memorising).""")

code("""data = torch.tensor(encode(text), dtype=torch.long)
print(data.shape, data.dtype)
print(data[:50])   # the first 50 chars, now as numbers

n = int(0.9 * len(data))
train_data = data[:n]
val_data = data[n:]
print("train:", len(train_data), " val:", len(val_data))""")

md("""## 4. Context: what does "predict the next token" actually mean?

We never feed the whole 1M-character book at once. We feed **chunks** of a fixed
length called `block_size` (the context window). Within one chunk of length 8 there
are actually 8 training examples packed in — predict char 2 from char 1, char 3 from
chars 1-2, and so on. This teaches the model to handle contexts of every length from
1 up to `block_size`.""")

code("""block_size = 8
x = train_data[:block_size]
y = train_data[1:block_size + 1]   # y is x shifted left by one
for t in range(block_size):
    context = x[:t + 1]
    target = y[t]
    print(f"when input is {context.tolist()} -> target is {target}")""")

md("""## 5. Batching

GPUs are happiest doing many things in parallel. So instead of one chunk, we grab a
**batch** of several random chunks and process them at once. Two key numbers:

- `batch_size` = how many independent chunks we stack (parallelism).
- `block_size` = how long each chunk is (context length).

`get_batch` picks random starting points and returns:
- `x`: the inputs, shape `(batch_size, block_size)`
- `y`: the targets (each input shifted by one), same shape.

`.to(device)` moves the tensors onto the GPU.""")

code("""torch.manual_seed(1337)
batch_size = 4
block_size = 8

def get_batch(split):
    d = train_data if split == "train" else val_data
    ix = torch.randint(len(d) - block_size, (batch_size,))   # random start indices
    x = torch.stack([d[i:i + block_size] for i in ix])
    y = torch.stack([d[i + 1:i + block_size + 1] for i in ix])
    return x.to(device), y.to(device)

xb, yb = get_batch("train")
print("inputs ", xb.shape)
print(xb)
print("targets", yb.shape)
print(yb)""")

md("""## 6. The Bigram model — our baseline

Simplest possible language model: to predict the next character, look **only** at the
current character. We do this with an `nn.Embedding` table of shape
`(vocab_size, vocab_size)`. Row `i` is directly the scores ("logits") for what comes
after token `i`.

Key concepts in this cell:
- **Logits**: raw, unnormalised scores for each possible next token.
- **Cross-entropy loss**: measures how surprised the model is by the true next token.
  Lower = better. With `vocab_size` random guesses we expect loss ~ `-ln(1/vocab_size)`.""")

code("""class BigramLanguageModel(nn.Module):
    def __init__(self, vocab_size):
        super().__init__()
        # each token directly reads the next-token logits from this lookup table
        self.token_embedding_table = nn.Embedding(vocab_size, vocab_size)

    def forward(self, idx, targets=None):
        logits = self.token_embedding_table(idx)   # (B, T, vocab_size)
        if targets is None:
            return logits, None
        # cross_entropy wants (N, C), so flatten the batch & time dims together
        B, T, C = logits.shape
        loss = F.cross_entropy(logits.view(B * T, C), targets.view(B * T))
        return logits, loss

    def generate(self, idx, max_new_tokens):
        # idx is (B, T) of current context; extend it one token at a time
        for _ in range(max_new_tokens):
            logits, _ = self(idx)
            logits = logits[:, -1, :]                 # only the last step -> (B, C)
            probs = F.softmax(logits, dim=-1)         # scores -> probabilities
            idx_next = torch.multinomial(probs, num_samples=1)  # sample one token
            idx = torch.cat((idx, idx_next), dim=1)
        return idx

model = BigramLanguageModel(vocab_size).to(device)
logits, loss = model(xb, yb)
expected = -torch.log(torch.tensor(1.0 / vocab_size)).item()
print("logits shape:", logits.shape)
print(f"initial loss: {loss.item():.4f}   (random guessing ~ {expected:.4f})")""")

md("""## 7. Generate *before* training

Right now the weights are random, so this will be pure gibberish. We start generation
from a single newline character (token 0). This is our "before" picture.""")

code("""context = torch.zeros((1, 1), dtype=torch.long, device=device)
print(decode(model.generate(context, max_new_tokens=200)[0].tolist()))""")

md("""## 8. Train it

The training loop is the same four steps forever:
1. **forward** — get predictions + loss
2. **zero_grad** — clear old gradients
3. **backward** — backpropagation computes how each weight affected the loss
4. **step** — the optimizer (AdamW) nudges weights to reduce the loss

> Analogy: loss is "how wrong we are." Backprop tells each knob which way to turn.
> The optimizer turns all the knobs a little. Repeat thousands of times.""")

code("""optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
batch_size = 32

for step in range(10000):
    xb, yb = get_batch("train")
    logits, loss = model(xb, yb)
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    optimizer.step()
    if step % 1000 == 0:
        print(f"step {step:5d}: loss {loss.item():.4f}")

print("final loss:", round(loss.item(), 4))""")

md("""## 9. Generate *after* training

Still not Shakespeare — a bigram model only ever looks at one character back, so it
can't form real words. But it should now produce plausible letter/space patterns
instead of random noise. That jump is the model *learning*.

The fix for "only one character back" is **self-attention** — letting each token look
at all previous tokens. That's the heart of the Transformer, and it's our next step.""")

code("""context = torch.zeros((1, 1), dtype=torch.long, device=device)
print(decode(model.generate(context, max_new_tokens=400)[0].tolist()))""")

nb = {
    "cells": cells,
    "metadata": {
        "kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"},
        "language_info": {"name": "python", "version": "3.12"},
    },
    "nbformat": 4,
    "nbformat_minor": 5,
}

# normalise source: nbformat stores source as a list of lines, each (except last) ending in \n
for c in nb["cells"]:
    s = c["source"]
    lines = s.split("\n")
    c["source"] = [ln + "\n" for ln in lines[:-1]] + [lines[-1]]

with open(os.path.join(HERE, "01_build_gpt.ipynb"), "w", encoding="utf-8") as f:
    json.dump(nb, f, indent=1)

print("wrote 01_build_gpt.ipynb with", len(cells), "cells")