# 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 numpy as np
from pathlib import Path
import matplotlib.pyplot as plt
np.random.seed(1337)

# --- setup from notebook 02 (tokenizer + data) ---
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 ids: "".join(itos[i] for i in ids)
data = np.array(encode(text), dtype=np.int64)
def softmax(x, axis=-1):
    x = x - x.max(axis=axis, keepdims=True); e = np.exp(x); return e / e.sum(axis=axis, keepdims=True)
print("vocab:", vocab_size, " dataset:", len(data))

W = np.random.randn(vocab_size, vocab_size) * 0.01    # the only parameters

def forward(x):                 # x: (B,) current-char ids
    logits = W[x]               # (B, vocab) - selecting rows IS the matmul with one-hots
    return logits

def loss_fn(logits, y):         # cross-entropy of the true next chars y
    p = softmax(logits)
    return -np.mean(np.log(p[np.arange(len(y)), y] + 1e-9))

# sanity: an untrained model should have loss ~ ln(vocab_size)
xb = data[:1000]; yb = data[1:1001]
print("initial loss:", round(loss_fn(forward(xb), yb), 4), " | ln(vocab) =", round(np.log(vocab_size), 4))

# The entire model, made visible.
print("W.shape:", W.shape, " -> rows = current char, cols = next char\n")

print("top-left corner W[:4, :6] (raw logits -- all tiny, near zero):")
print(np.round(W[:4, :6], 4))

# One row is the model's whole opinion about "what follows this character".
# Fresh, that row is near-flat, so softmax turns it into a near-uniform guess.
row = softmax(W[stoi["q"]])
best = np.argsort(row)[::-1][:3]
print("\nafter 'q', the untrained model's top-3 next characters:")
for j in best:
    print(f"   {itos[j]!r:>4}  p = {row[j]:.4f}")
print(f"a flat, no-opinion guess is 1/65 = {1/vocab_size:.4f}  <- essentially the same")

def backward(x, y):
    B = len(x)
    p = softmax(W[x])                 # (B, vocab)
    dlogits = p.copy()
    dlogits[np.arange(B), y] -= 1     # probs - one_hot(target)
    dlogits /= B                      # average over the batch
    dW = np.zeros_like(W)
    np.add.at(dW, x, dlogits)         # route each row's gradient back to W[x]
    return dW

# gradient check: compare our analytic dW to a numerical estimate on a few entries
xb, yb = data[:256], data[1:257]
dW = backward(xb, yb)
def num_grad(i, j, h=1e-4):
    global W
    W[i, j] += h; lp = loss_fn(forward(xb), yb)
    W[i, j] -= 2*h; lm = loss_fn(forward(xb), yb)
    W[i, j] += h
    return (lp - lm) / (2*h)
for (i, j) in [(10, 20), (5, 5), (40, 1)]:
    print(f"  W[{i},{j}]  analytic={dW[i,j]: .6f}  numerical={num_grad(i,j): .6f}")
print("they match -> our hand-derived gradient is correct")

# Rebuild the gradient for ONE example and look at the signal, number by number.
xi, yi = int(xb[0]), int(yb[0])            # current char, true next char
probs = softmax(W[xi])                     # the model's guess for this row
g = probs.copy(); g[yi] -= 1               # probs - one_hot(target): the whole signal

print(f"current char {itos[xi]!r}  ->  true next char {itos[yi]!r}")
print(f"prob the model put on the correct answer : {probs[yi]:.4f}   (we want this near 1)")
print(f"gradient at the correct column           : {g[yi]:+.4f}   (negative -> push that logit UP)")
other = g[np.arange(vocab_size) != yi]
print(f"gradient at each of the 64 other columns : ~{other.mean():+.4f}    (positive -> push them DOWN)")

plt.figure(figsize=(6, 2.4))
colors = np.where(np.arange(vocab_size) == yi, "#e5484d", "#5b5bd6")
plt.bar(range(vocab_size), g, color=colors)
plt.axhline(0, color="k", lw=.5)
plt.title(f"gradient on the logits for one '{itos[xi]}' -> '{itos[yi]}' example")
plt.xlabel("next-char id"); plt.ylabel("d loss / d logit"); plt.show()

# One real pass of the loop, narrated. We draw a batch the same way the loop does
# (a private RNG, so the seeded loop below still reproduces its numbers), then
# read off what each step produces.
rng = np.random.default_rng(0)
B = 1024
ix = rng.integers(0, len(data) - 1, size=B)
xb1, yb1 = data[ix], data[ix + 1]              # step 1: B (current -> next) pairs

print(f"step 1  a batch of {B} (current -> next) pairs; the first 6:")
for cx, cy in zip(xb1[:6], yb1[:6]):
    print(f"          id {cx:>2} {itos[cx]!r:>4}   ->   id {cy:>2} {itos[cy]!r:>4}")

logits = forward(xb1)                          # step 2: forward
L = loss_fn(logits, yb1)                       # step 3: loss
dW = backward(xb1, yb1)                         # step 4: backward
rows_touched = int((dW != 0).any(axis=1).sum())
print(f"\nstep 2  forward : logits.shape = {logits.shape}  ({B} rows, a 65-long score vector each)")
print(f"step 3  loss    : {L:.4f} nats   (near ln 65 = 4.17 -- the model is still untrained)")
print(f"step 4  backward: dW.shape = {dW.shape}; {rows_touched} of 65 rows are nonzero")
print(f"                  -> a character only gets a nudge if it appeared in the batch")

i, jt = int(xb1[0]), int(yb1[0])               # step 5: the update rule, on one weight
print(f"step 5  update  : W -= lr * dW.  dW[{i},{jt}] = {dW[i, jt]:+.5f} for the seen "
      f"'{itos[i]}'->'{itos[jt]}' pair,")
print(f"                  so W[{i},{jt}] steps {W[i, jt]:+.4f} -> {W[i, jt] - 30.0 * dW[i, jt]:+.4f}   (lr=30)")
print("step 6  repeat this thousands of times -- that is the whole cell below.")

W = np.random.randn(vocab_size, vocab_size) * 0.01   # fresh start
lr, steps, batch_size = 30.0, 6000, 1024
losses = []
for step in range(steps):
    ix = np.random.randint(0, len(data) - 1, size=batch_size)
    x, y = data[ix], data[ix + 1]
    logits = forward(x)
    losses.append(loss_fn(logits, y))
    W -= lr * backward(x, y)

# one batch is a noisy reading (easy and hard stretches of text exist),
# so report an average over the last 200 steps
print("final train loss (avg of last 200 steps):", round(float(np.mean(losses[-200:])), 4))

plt.figure(figsize=(6, 3))
plt.plot(losses); plt.title("bigram training loss"); plt.xlabel("step"); plt.ylabel("loss"); plt.show()

# Same table as before, now trained. Turn each row into next-char probabilities and look.
P = softmax(W, axis=1)

plt.figure(figsize=(7, 6))
plt.imshow(P, cmap="magma", aspect="auto")
plt.colorbar(label="P(next char | current char)")
plt.xticks(range(vocab_size), [itos[i] for i in range(vocab_size)], fontsize=5)
plt.yticks(range(vocab_size), [itos[i] for i in range(vocab_size)], fontsize=5)
plt.xlabel("next character"); plt.ylabel("current character")
plt.title("what the trained bigram learned"); plt.show()

# Read a few rows back in plain English -- the model was never told any of this.
for c in ["q", ".", " "]:
    r = P[stoi[c]]; j = int(r.argmax())
    name = "space" if c == " " else repr(c)
    print(f"after {name:>7}: most likely next char is {itos[j]!r:>4}  (p = {r[j]:.2f})")

# Count every adjacent pair of characters in the data, turn each row of counts into
# probabilities, and measure the loss of those PERFECT bigram probabilities.
# No model that sees only one character of context can ever beat this number --
# it is the built-in uncertainty of "next char given current char" on this text.
counts = np.zeros((vocab_size, vocab_size))
np.add.at(counts, (data[:-1], data[1:]), 1)
p = counts / np.maximum(counts.sum(axis=1, keepdims=True), 1)
row_weight = counts.sum(axis=1) / counts.sum()            # how often each char occurs
with np.errstate(divide="ignore", invalid="ignore"):
    row_entropy = -np.nansum(np.where(p > 0, p * np.log(p), 0.0), axis=1)
optimum = float(np.nansum(row_weight * row_entropy))
print(f"bigram optimum on this text: {optimum:.4f} nats")
print("our trained model sits essentially at this ceiling -- the only way past it is more context")

# Put the training curve and the two reference levels on one axis.
plt.figure(figsize=(6.5, 3.2))
plt.plot(losses, color="#5b5bd6", label="training loss")
plt.axhline(np.log(vocab_size), color="#999", ls=":", label=f"random start = {np.log(vocab_size):.2f}")
plt.axhline(optimum, color="#e5484d", ls="--", label=f"bigram ceiling = {optimum:.2f}")
plt.xlabel("step"); plt.ylabel("loss"); plt.legend()
plt.title("the loss falls from random and settles on the ceiling"); plt.show()

gap = float(np.mean(losses[-200:])) - optimum
print(f"gap left between our model and the perfect bigram: {gap:+.4f} nats")
print("essentially zero -- one character of context has nothing left to give")

def generate(start, n=400, seed=0):
    rng = np.random.default_rng(seed)
    idx = stoi[start]; out = [idx]
    for _ in range(n):
        p = softmax(W[idx])
        idx = int(rng.choice(vocab_size, p=p))
        out.append(idx)
    return decode(out)

print(generate("T", 400))