"""Self-attention, worked as Chapter 16 works it.

Five things, all reproducing the chapter's tables:

    one query          the lecture's three-key example, every number
    when it selects    a query that actually picks a key, and one that cannot
    why root d_k       what the scaling stops the softmax from doing
    the causal mask    mask then softmax, and why never the other way round
    the parameter count where a transformer layer's numbers sit

    python3 attention.py            # all five
    python3 attention.py --one      # the lecture example
    python3 attention.py --select   # attention that discriminates
    python3 attention.py --scale    # the square root, tested
    python3 attention.py --mask     # causal masking
    python3 attention.py --params   # one layer, counted

Install:  pip install numpy
"""

import argparse
import math

try:
    import numpy as np
except ImportError:
    raise SystemExit("this one needs numpy:  pip install numpy")


def softmax(z):
    e = np.exp(z - np.max(z, axis=-1, keepdims=True))
    return e / e.sum(axis=-1, keepdims=True)


def entropy(p):
    """In bits. Maximum is log2(n) for a uniform distribution."""
    p = p[p > 0]
    return float(-(p * np.log2(p)).sum())


# ---------------------------------------------------------------- one query

def show_one():
    q = np.array([0.1, 0.2, 0.3])
    keys = {"is": [0.4, 0.5, 0.6], "the": [0.2, 0.1, 0.3],
            "best": [0.6, 0.7, 0.8]}
    K = np.array(list(keys.values()))
    V = K.copy()                       # values equal keys, as the lecture has it
    d = q.shape[0]

    print("\nSTEP 1  one query against three keys\n")
    print(f"  query 'who'  q = {q}")
    for w, k in keys.items():
        label = f"key '{w}'"
        print(f"  {label:<13} k = {np.array(k)}")
    print(f"\n  Values equal keys here, to keep the arithmetic short.")
    print(f"  d_k = {d}, so the scale is sqrt({d}) = {math.sqrt(d):.4f}\n")

    raw = K @ q
    scaled = raw / math.sqrt(d)
    ex = np.exp(scaled)
    alpha = ex / ex.sum()

    print(f"  {'key':<7}{'q . k':>9}{'/ sqrt(d)':>12}{'exp':>9}{'alpha':>9}")
    print("  " + "-" * 48)
    for w, r, s, e, a in zip(keys, raw, scaled, ex, alpha):
        print(f"  {w:<7}{r:>9.4f}{s:>12.4f}{e:>9.4f}{a:>9.4f}")
    print("  " + "-" * 48)
    print(f"  {'sum':<7}{'':>9}{'':>12}{ex.sum():>9.4f}{alpha.sum():>9.4f}")

    z = alpha @ V
    print(f"\n  output  zeta = sum of alpha_i v_i")
    for w, a, v in zip(keys, alpha, V):
        print(f"    {a:.4f} x {v}")
    print(f"    = ({z[0]:.4f}, {z[1]:.4f}, {z[2]:.4f})")

    print(f"\n  Now look at the alpha column: {alpha[0]:.2f},"
          f" {alpha[1]:.2f}, {alpha[2]:.2f}.")
    print(f"  That is almost uniform. Its entropy is {entropy(alpha):.4f} bits")
    print(f"  against a maximum of {math.log2(3):.4f}.")
    print(f"\n  So this query attends to everything about equally, and the")
    print(f"  output is close to the plain average of the values. The")
    print(f"  mechanism ran correctly and selected nothing, because these")
    print(f"  three keys all point the same way. Step 2 fixes that.")


# ------------------------------------------------------------ when it selects

def show_select():
    print("\n\nSTEP 2  attention that actually discriminates\n")
    print("  Same machinery, keys that disagree. Three dimensions standing")
    print("  for (animal, vehicle, verb).\n")

    keys = {"dog": [0.9, 0.0, 0.1], "truck": [0.0, 0.9, 0.1],
            "barked": [0.1, 0.0, 0.9]}
    K = np.array(list(keys.values()))
    V = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=float)
    d = 3

    for name, q in [("an animal query", np.array([5.0, 0.0, 0.0])),
                    ("a verb query", np.array([0.0, 0.0, 5.0])),
                    ("an undecided query", np.array([2.0, 2.0, 2.0]))]:
        alpha = softmax(K @ q / math.sqrt(d))
        z = alpha @ V
        cells = "  ".join(f"{w} {a:.4f}" for w, a in zip(keys, alpha))
        print(f"  {name:<20}q = {q}")
        print(f"    alpha  {cells}")
        print(f"    output ({z[0]:.4f}, {z[1]:.4f}, {z[2]:.4f})"
              f"   entropy {entropy(alpha):.4f} bits\n")

    print(f"  The first two queries commit. The third cannot, and its")
    print(f"  entropy sits at the {math.log2(3):.4f} bit maximum for three keys.")
    print("\n  Low entropy means the head has made a choice. High entropy")
    print("  means it is averaging. Both are legitimate, and reading which")
    print("  one a head is doing is how attention maps get interpreted.")


# ------------------------------------------------------------- the scaling

def show_scale():
    print("\n\nSTEP 3  why divide by the square root of d_k\n")
    print("  Take random query and key vectors with unit variance entries.")
    print("  Their dot product is a sum of d terms, so its standard")
    print("  deviation grows like sqrt(d). Left alone, the scores get")
    print("  larger and larger as the model gets wider.\n")

    rng = np.random.default_rng(0)
    n = 8
    print(f"  {'d_k':>6}{'sd of q.k':>12}{'entropy raw':>14}"
          f"{'entropy scaled':>16}{'max alpha raw':>15}")
    print("  " + "-" * 66)
    for d in (4, 16, 64, 256, 1024):
        sd, e_raw, e_scaled, top = [], [], [], []
        for _ in range(200):
            q = rng.normal(0, 1, d)
            K = rng.normal(0, 1, (n, d))
            s = K @ q
            sd.append(s.std())
            a_raw = softmax(s)
            a_sc = softmax(s / math.sqrt(d))
            e_raw.append(entropy(a_raw))
            e_scaled.append(entropy(a_sc))
            top.append(a_raw.max())
        print(f"  {d:>6}{np.mean(sd):>12.3f}{np.mean(e_raw):>14.4f}"
              f"{np.mean(e_scaled):>16.4f}{np.mean(top):>15.4f}")

    print(f"\n  Maximum possible entropy over {n} keys is"
          f" {math.log2(n):.4f} bits.")
    print("\n  Read the raw column. By d = 1024 the distribution has")
    print("  collapsed onto a single key, before any training has happened.")
    print("  A softmax that saturated is a softmax with no gradient, so the")
    print("  head cannot learn its way out.")
    print("\n  The scaled column barely moves. Dividing by sqrt(d_k) cancels")
    print("  exactly the growth the second column shows, which is why the")
    print("  scale is a square root and not something tuned.")


# ------------------------------------------------------------- causal mask

def show_mask():
    print("\n\nSTEP 4  the causal mask, and the order of operations\n")
    print("  A language model must not read the future. When scoring")
    print("  position i, only positions j <= i may contribute.\n")
    print("  Set the blocked scores to minus infinity BEFORE the softmax:\n")
    print("    S_ij = (Q K^T)_ij / sqrt(d_k)   if j <= i,   else -inf\n")

    S = np.array([[2.0, 1.0, 0.5, 1.5],
                  [0.5, 2.0, 1.0, 0.5],
                  [1.0, 0.5, 2.0, 1.0],
                  [1.5, 1.0, 0.5, 2.0]])
    n = len(S)
    tokens = ["the", "dog", "barked", "loudly"]

    print(f"  raw scores S (rows are queries, columns are keys):\n")
    print("  " + " " * 10 + "".join(f"{t:>10}" for t in tokens))
    for i, t in enumerate(tokens):
        print(f"  {t:<10}" + "".join(f"{v:>10.2f}" for v in S[i]))

    mask = np.tril(np.ones((n, n), dtype=bool))
    Sm = np.where(mask, S, -np.inf)
    A = softmax(Sm)

    print(f"\n  after masking then softmax:\n")
    print("  " + " " * 10 + "".join(f"{t:>10}" for t in tokens) + f"{'sum':>8}")
    for i, t in enumerate(tokens):
        cells = "".join("         ." if not mask[i][j] else f"{A[i][j]:>10.4f}"
                        for j in range(n))
        print(f"  {t:<10}{cells}{A[i].sum():>8.4f}")

    print("\n  Every row still sums to 1. That is the whole reason the mask")
    print("  goes before the softmax rather than after.\n")

    wrong = softmax(S) * mask
    print("  Compare softmax first, then zero the blocked cells:\n")
    print("  " + " " * 10 + "".join(f"{t:>10}" for t in tokens) + f"{'sum':>8}")
    for i, t in enumerate(tokens):
        cells = "".join("         ." if not mask[i][j] else f"{wrong[i][j]:>10.4f}"
                        for j in range(n))
        print(f"  {t:<10}{cells}{wrong[i].sum():>8.4f}")

    print(f"\n  Row 1 now sums to {wrong[0].sum():.4f} instead of 1. The mass")
    print("  that belonged to the blocked positions has simply been thrown")
    print("  away rather than redistributed.")
    print("\n  Masking first lets the surviving positions inherit that mass.")
    print("  Masking second silently scales the whole row down, and the")
    print("  earlier the token, the worse it gets.")


# --------------------------------------------------------- the parameter count

def show_params(V=50000, D=512, L=6, h=8, H=2048):
    d_k = d_v = D // h
    print(f"\n\nSTEP 5  a transformer layer, counted\n")
    print(f"  |V| = {V}, D = {D}, layers L = {L}, heads h = {h},")
    print(f"  d_k = d_v = D/h = {d_k}, feed-forward H = {H}\n")

    attn = 4 * D * D                       # W^Q, W^K, W^V, W^O
    ffn = 2 * D * H + H + D
    ln = 4 * D                             # two layer norms, scale and shift
    layer = attn + ffn + ln
    emb = V * D

    print(f"  {'component':<26}{'formula':<22}{'parameters':>14}")
    print("  " + "-" * 64)
    print(f"  {'attention Q,K,V,O':<26}{'4 D^2':<22}{attn:>14,}")
    print(f"  {'feed-forward':<26}{'2 D H + H + D':<22}{ffn:>14,}")
    print(f"  {'layer norms':<26}{'4 D':<22}{ln:>14,}")
    print("  " + "-" * 64)
    print(f"  {'one layer':<26}{'':<22}{layer:>14,}")
    print(f"  {'all ' + str(L) + ' layers':<26}{'L x layer':<22}"
          f"{L * layer:>14,}")
    print(f"  {'embeddings':<26}{'V D':<22}{emb:>14,}")
    print("  " + "-" * 64)
    print(f"  {'total':<48}{L * layer + emb:>14,}")

    print(f"\n  Two things to notice.")
    print(f"\n  The heads are free. Splitting D into {h} heads of {d_k} costs")
    print(f"  nothing, because h x d_k = D. Multi-head attention is a")
    print(f"  reshape, not an extra parameter budget.")
    print(f"\n  The feed-forward block is larger than the attention block,")
    print(f"  {ffn:,} against {attn:,}. Attention gets the attention, and")
    print(f"  most of the parameters sit next door.")

    print(f"\n\n  WHAT THE SEQUENCE LENGTH COSTS\n")
    print(f"  No parameter count above mentions sequence length. The")
    print(f"  computation does: every position attends to every position.\n")
    print(f"  {'tokens n':>10}{'n^2 scores per head':>22}{'vs n = 512':>13}")
    print("  " + "-" * 46)
    for nt in (128, 512, 2048, 8192, 32768):
        print(f"  {nt:>10}{nt * nt:>22,}{nt * nt / 512 ** 2:>12.2f}x")
    print(f"\n  Quadruple the context and the attention cost goes up")
    print(f"  sixteen-fold. That is the wall the efficiency chapter climbs.")


def main():
    ap = argparse.ArgumentParser(description=__doc__)
    for f in ("one", "select", "scale", "mask", "params"):
        ap.add_argument(f"--{f}", action="store_true")
    a = ap.parse_args()
    picked = a.one or a.select or a.scale or a.mask or a.params
    if a.one or not picked:
        show_one()
    if a.select or not picked:
        show_select()
    if a.scale or not picked:
        show_scale()
    if a.mask or not picked:
        show_mask()
    if a.params or not picked:
        show_params()
    print()


if __name__ == "__main__":
    main()