"""
llm.py — numerical companion to the glossary entry
'large language model (LLM)'.

One block per paragraph of the entry (marked [P...]): each block verifies
numerically what the corresponding statement asserts. Self-contained
(numpy/matplotlib only), fixed seed. (The entry's scale claim —
billions of parameters — is illustrated in miniature: the mechanisms,
not the size.)

Blocks
------
[P-selfsup]   Self-supervised construction of labeled data points from
              raw text: masking words turns an unannotated corpus into
              (context features, masked-word label) pairs — one per
              position, with zero human labeling effort.
[P-train]     Training via ERM on these pairs: a small next-token model
              (embedding + softmax over the vocabulary) trained by
              gradient descent on the training loss (the negative log probability of the correct next token) drives
              training loss down and beats the uniform-guess baseline.
[P-nexttoken] A trained LLM maps an input token sequence to a
              probability distribution over the next token: outputs are
              nonnegative, sum to one, and the model assigns the
              highest probability to continuations seen in the corpus;
              sampling from the distribution generates text.

Outputs
-------
llm.png : preview figure (checking only).

Data generated by pythondemos/llm.py.
"""

import numpy as np
import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt

rng = np.random.default_rng(42)
report = []


def check(name, ok):
    report.append((name, bool(ok)))
    print(f"  [{'ok' if ok else 'FAIL'}] {name}")


corpus = ("all human beings are born free and equal in dignity and "
          "rights all human beings are endowed with reason and "
          "conscience").split()
vocab = sorted(set(corpus))
V = len(vocab)
tok = {w: i for i, w in enumerate(vocab)}
ids = np.array([tok[w] for w in corpus])

# -------------------------------------------------------- [P-selfsup]
print("[P-selfsup] masked words become labels, contexts become features")
pairs = [(ids[t], ids[t + 1]) for t in range(len(ids) - 1)]
check("one labeled pair per corpus position (no human annotation)",
      len(pairs) == len(corpus) - 1)
check("labels are drawn from the text itself",
      all(0 <= y < V for _, y in pairs))

# ---------------------------------------------------------- [P-train]
print("[P-train] ERM on the constructed pairs")
d_emb = 8
E = 0.1 * rng.normal(size=(V, d_emb))              # embeddings
U = 0.1 * rng.normal(size=(d_emb, V))              # unembedding
def forward(x_ids):
    logits = E[x_ids] @ U
    p = np.exp(logits - logits.max(1, keepdims=True))
    return p / p.sum(1, keepdims=True)
xs = np.array([x for x, _ in pairs])
ys = np.array([y for _, y in pairs])
def xent():
    return -np.mean(np.log(forward(xs)[np.arange(len(ys)), ys] + 1e-12))
loss0 = xent()
for _ in range(800):                               # gradient descent
    P = forward(xs)
    G = P.copy(); G[np.arange(len(ys)), ys] -= 1.0
    gU = E[xs].T @ G / len(ys)
    gE = np.zeros_like(E)
    np.add.at(gE, xs, G @ U.T / len(ys))
    U -= 2.0 * gU; E -= 2.0 * gE
loss1 = xent()
print(f"    training loss: init {loss0:.2f} -> trained {loss1:.2f} "
      f"(uniform baseline {np.log(V):.2f})")
check("training reduces the training loss", loss1 < loss0)
check("the trained model beats the uniform-guess baseline log|V|",
      loss1 < np.log(V) - 0.5)

# ------------------------------------------------------ [P-nexttoken]
print("[P-nexttoken] input sequence -> distribution over the next token")
p_next = forward(np.array([tok["human"]]))[0]
check("the output is a probability distribution (nonneg, sums to 1)",
      np.all(p_next >= 0) and np.isclose(p_next.sum(), 1.0))
check("'human' is followed by 'beings' in the corpus — and gets the "
      "highest next-token probability",
      vocab[int(np.argmax(p_next))] == "beings")
gen = [tok["all"]]
for _ in range(5):                                 # sample a continuation
    gen.append(int(rng.choice(V, p=forward(np.array([gen[-1]]))[0])))
check("sampling from the distributions generates a token sequence",
      len(gen) == 6 and all(0 <= g < V for g in gen))
print("    generated:", " ".join(vocab[g] for g in gen))

# ------------------------------------------------------------ preview
fig, ax = plt.subplots(figsize=(5.6, 3.0))
ax.bar(range(V), p_next)
ax.set_xticks(range(V))
ax.set_xticklabels(vocab, rotation=90, fontsize=6)
ax.set_ylabel("P(next token | 'human')")
ax.set_title("[P-nexttoken] next-token distribution")
fig.tight_layout()
fig.savefig("llm.png", dpi=110)
print(f"\n{sum(ok for _, ok in report)}/{len(report)} checks passed")
assert all(ok for _, ok in report)