"""
Lecture 20 - Scaled Dot-Product Self-Attention from scratch (NumPy).

Mirrors the lecture's Section 4 walkthrough:
  Step 1: token embeddings (X)
  Step 2: project to Query, Key, Value   Q = X.Wq   K = X.Wk   V = X.Wv
  Step 3: scores = Q K^T / sqrt(d_k), then row-wise softmax
  Step 4: output Z = Attention @ V  (a context-blended, "contextual" embedding)

The numbers below are made up (seeded random) -- NOT the instructor's
truncated 768-dim slide numbers, which can't be reproduced exactly from
a small demo. The qualitative shape should match the lecture's heatmap:
rows sum to 1, some rows peaked (near-certain self-attention), some spread
across several tokens.
"""
import numpy as np

np.random.seed(20)


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


tokens = ["The", "dog", "ran", "fast"]
n_tokens, d_model, d_k = 4, 8, 8

# ---- Step 1: token embeddings X (n_tokens x d_model) ----
X = np.random.randn(n_tokens, d_model)

# ---- Step 2: learned projection matrices -> Q, K, V ----
Wq = np.random.randn(d_model, d_k) * 0.5
Wk = np.random.randn(d_model, d_k) * 0.5
Wv = np.random.randn(d_model, d_k) * 0.5

Q = X @ Wq
K = X @ Wk
V = X @ Wv

# ---- Step 3: raw scores QK^T, scale by sqrt(d_k), softmax per row ----
scores = Q @ K.T
scaled_scores = scores / np.sqrt(d_k)
attention = softmax(scaled_scores, axis=1)

print("Attention matrix (each row is a probability distribution, sums to 1):")
for tok, row in zip(tokens, attention):
    row_str = "  ".join(f"{v:.3f}" for v in row)
    print(f"  {tok:>5}: {row_str}   sum={row.sum():.3f}")

# ---- Step 4: weighted sum of Values -> contextual output Z ----
Z = attention @ V

print("\nOutput Z (contextual embeddings, one row per token, first 4 dims shown):")
for tok, row in zip(tokens, Z):
    row_str = "  ".join(f"{v:.3f}" for v in row[:4])
    print(f"  {tok:>5}: {row_str} ...")