# Colab setup -- fetch the files this notebook needs.
# (Does nothing when run locally in the course folder.)
import os, urllib.request
BASE = ("https://raw.githubusercontent.com/waze"
        "emlabs/llm-book-code/main/")
for f in ['data/input.txt']:
    if not os.path.exists(f):
        d = os.path.dirname(f)
        if d: os.makedirs(d, exist_ok=True)
        urllib.request.urlretrieve(BASE + f, f)
        print("downloaded", f)

import torch
import torch.nn as nn
from torch.nn import functional as F
torch.manual_seed(1337)

device = "mps" if torch.backends.mps.is_available() else ("cuda" if torch.cuda.is_available() else "cpu")
print("torch", torch.__version__, "| device:", device)

a = torch.tensor(2.0, requires_grad=True)
b = torch.tensor(-3.0, requires_grad=True)
c = torch.tensor(0.5, requires_grad=True)

L = (a*b + b**2) * torch.tanh(c)
L.backward()                       # autograd walks the graph for us
print("a.grad = %.5f   b.grad = %.5f   c.grad = %.5f" % (a.grad, b.grad, c.grad))
print("(notebook 04 by hand:  a=-1.38635  b=-1.84847  c=2.35934)")

# PyTorch recorded every operation in L.grad_fn. Walk it, exactly like notebook 04.
leaf = {id(a): "a", id(b): "b", id(c): "c"}
L = (a*b + b**2) * torch.tanh(c)        # rebuild so the graph is intact to walk

def show(fn, depth=0):
    if fn is None:
        return
    name = type(fn).__name__
    if name == "AccumulateGrad":                    # a leaf input (a, b, or c)
        print("   " * depth + f"leaf  {leaf.get(id(fn.variable), '?')}")
    else:
        print("   " * depth + name)                 # e.g. MulBackward0, TanhBackward0
        for nxt, _ in fn.next_functions:
            show(nxt, depth + 1)

show(L.grad_fn)
print("\nthe same graph as notebook 04: two *, one +, one **2, one tanh, leaves a/b/c")

from pathlib import Path
# --- data + tokenizer (notebook 02) ---
text = Path("data/input.txt").read_text()
chars = sorted(set(text)); vocab_size = len(chars)
stoi = {c:i for i,c in enumerate(chars)}; itos = {i:c for c,i in stoi.items()}
encode = lambda s: [stoi[c] for c in s]; decode = lambda l: "".join(itos[i] for i in l)
data = torch.tensor(encode(text), dtype=torch.long)
n = int(0.9*len(data)); train_data, val_data = data[:n], data[n:]

block_size, batch_size = 8, 32
def get_batch(split):
    d = train_data if split=="train" else val_data
    ix = torch.randint(len(d)-block_size, (batch_size,))
    x = torch.stack([d[i:i+block_size] for i in ix])
    y = torch.stack([d[i+1:i+1+block_size] for i in ix])
    return x.to(device), y.to(device)
print("vocab:", vocab_size, "| one batch:", tuple(get_batch("train")[0].shape))

class BigramLM(nn.Module):
    def __init__(self, vocab_size):
        super().__init__()
        # each token id directly indexes a row of next-token logits
        self.token_emb = nn.Embedding(vocab_size, vocab_size)
    def forward(self, idx, targets=None):
        logits = self.token_emb(idx)                 # (B,T,vocab)
        loss = None
        if targets is not None:
            B,T,V = logits.shape
            loss = F.cross_entropy(logits.view(B*T, V), targets.view(B*T))
        return logits, loss

model = BigramLM(vocab_size).to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-2)

for step in range(4000):
    x, y = get_batch("train")
    _, loss = model(x, y)
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    optimizer.step()

# a single batch is a noisy reading, so evaluate properly over many batches
with torch.no_grad():
    eval_loss = torch.stack([model(*get_batch("train"))[1] for _ in range(200)]).mean()
print(f"final loss (avg over 200 batches): {eval_loss.item():.3f}")
print("right at the ~2.45 bigram ceiling we computed in notebook 03 -- same model, zero gradient code")

class Head(nn.Module):
    def __init__(self, n_embd, head_size, block_size):
        super().__init__()
        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)
        self.register_buffer("tril", torch.tril(torch.ones(block_size, block_size)))
    def forward(self, x):
        B,T,C = x.shape
        k, q = self.key(x), self.query(x)
        wei = q @ k.transpose(-2,-1) * k.shape[-1]**-0.5      # scaled affinities
        wei = wei.masked_fill(self.tril[:T,:T]==0, float('-inf'))  # causal mask
        wei = F.softmax(wei, dim=-1)
        return wei @ self.value(x)

x = torch.randn(4, 8, 32)
head = Head(n_embd=32, head_size=16, block_size=8)
out = head(x)
print("attention output:", tuple(out.shape))
# verify causality on the attention weights
with torch.no_grad():
    k,q = head.key(x), head.query(x)
    wei = (q @ k.transpose(-2,-1) * k.shape[-1]**-0.5).masked_fill(head.tril[:8,:8]==0, float('-inf')).softmax(-1)
print("causal (upper triangle zero):", torch.allclose(torch.triu(wei[0], 1), torch.zeros(8,8)))

# nn.Module auto-registers its weights (parameters) and its fixed helpers (buffers).
print("parameters (the optimizer trains these):")
for name, p in head.named_parameters():
    print(f"   {name:14s} {tuple(p.shape)}   {p.numel()} numbers")
print("buffers (fixed, moved to the device but never trained):")
for name, buf in head.named_buffers():
    print(f"   {name:14s} {tuple(buf.shape)}")
print(f"\ntotal trainable numbers in this head: {sum(p.numel() for p in head.parameters())}")
print("model.parameters() is exactly what you hand the optimizer -- no manual list to keep.")