"""Attention on one short sentence, step by step, in plain Python.

The sentence is "The cat sat because it was tired". Each word gets six numbers, and each head gets three small maps (query, key
and value). Every number is picked by hand so that the output is easy to read: nothing here was trained, and a real model learns
its numbers from data and uses thousands per word. The five steps below are the scaled dot-product attention of Vaswani et al.,
"Attention Is All You Need" (arXiv 1706.03762, section 3.2): for each word, score every word before it, hide the words after
it, turn the scores into weights that add up to 1, and take the weighted average of the earlier words' values.

    scores   = q . k / sqrt(d)            one per earlier word, d = how many numbers a query has
    weights  = softmax(scores)            positive, and they add up to 1
    output   = sum of weight * value      a blend of the earlier words' values

Python 3 standard library only. The widget (../widget.html) repeats these steps in JavaScript; test_widget.py runs both and
compares them, and test_attention.py checks this file against NumPy matrix operations.
"""
from __future__ import annotations

import math
import random

# --- the sentence, and six numbers per word: four about the word, then its place in the sentence (the cosine and sine of an angle) ---
TOKENS = ["The", "cat", "sat", "because", "it", "was", "tired"]
ABOUT = ["an animal", "pointing to something", "an action or state", "a small function word"]
CONTENT = [
    [0.0, 0.0, 0.0, 1.0],   # The
    [1.0, 0.0, 0.0, 0.0],   # cat      the only animal
    [0.0, 0.6, 1.0, 0.0],   # sat      an action that needs someone to do it
    [0.0, 0.0, 0.0, 1.0],   # because
    [0.0, 1.0, 0.0, 0.0],   # it       points to something, and carries nothing about what
    [0.0, 0.4, 1.0, 0.0],   # was
    [0.0, 0.8, 1.0, 0.0],   # tired    describes something
]
ANGLE = 0.9                  # radians between neighbouring positions


def position(p: int) -> list[float]:
    return [round(math.cos(ANGLE * p), 2), round(math.sin(ANGLE * p), 2)]


VECTORS = [CONTENT[p] + position(p) for p in range(len(TOKENS))]
D_MODEL = len(VECTORS[0])    # numbers per word: 6
D_K = 2                      # numbers per query and per key: 2
D_V = 4                      # numbers per value: 4


# --- three maps per head. A map is a table with D_MODEL rows and as many columns as the vector it makes. ---
def zeros(rows: int, cols: int) -> list[list[float]]:
    return [[0.0] * cols for _ in range(rows)]


def identity_on_content() -> list[list[float]]:
    value = zeros(D_MODEL, D_V)                  # the value keeps the four numbers about the word and drops the position
    for i in range(D_V):
        value[i][i] = 1.0
    return value


def refers_head(strength: float = 3.0) -> dict:
    """Looks for what a word points to: a word that points to something (column 1 of CONTENT) asks for an animal (column 0)."""
    query, key = zeros(D_MODEL, D_K), zeros(D_MODEL, D_K)
    query[1][0] = strength                       # "pointing to something" becomes the first number of the query
    key[0][0] = 1.0                              # "an animal" becomes the first number of the key
    return {"name": "refers", "wq": query, "wk": key, "wv": identity_on_content()}


def before_head(strength: float = 6.0) -> dict:
    """Looks for the word just before: the query is this word's position turned back by one step, so it matches the key of the word at p - 1."""
    c, s = round(strength * math.cos(ANGLE), 2), round(strength * math.sin(ANGLE), 2)
    query, key = zeros(D_MODEL, D_K), zeros(D_MODEL, D_K)
    query[4][0], query[4][1] = c, -s             # rows 4 and 5 are the position numbers
    query[5][0], query[5][1] = s, c
    key[4][0], key[5][1] = 1.0, 1.0              # the key is the position itself
    return {"name": "before", "wq": query, "wk": key, "wv": identity_on_content()}


def even_head() -> dict:
    """Asks for nothing: every query is all zeros, so every earlier word scores 0 and the weights come out equal."""
    return {"name": "even", "wq": zeros(D_MODEL, D_K), "wk": refers_head()["wk"], "wv": identity_on_content()}


HEADS = {"refers": refers_head(), "before": before_head(), "even": even_head()}
HEAD_LABELS = {"refers": "Refers to", "before": "Word before", "even": "No preference"}


# --- the five steps ---
def times(vector: list[float], table: list[list[float]]) -> list[float]:
    """A vector times a map: one output number per column."""
    return [sum(vector[r] * table[r][c] for r in range(len(vector))) for c in range(len(table[0]))]


def dot(a: list[float], b: list[float]) -> float:
    return sum(x * y for x, y in zip(a, b))


def scores_for(q: list[float], keys: list[list[float]], i: int, scale: bool = True) -> list[float]:
    """Step 2 and 3: q . k / sqrt(d) for every word, with -inf for the words after word i (the causal mask)."""
    divisor = math.sqrt(len(q)) if scale else 1.0
    return [dot(q, keys[j]) / divisor if j <= i else -math.inf for j in range(len(keys))]


def softmax(scores: list[float]) -> list[float]:
    """Step 4: positive weights that add up to 1. Subtracting the biggest score first changes nothing and avoids overflow; exp(-inf) is 0."""
    top = max(scores)
    sizes = [math.exp(s - top) for s in scores]
    total = sum(sizes)
    return [size / total for size in sizes]


def weighted_average(weights: list[float], values: list[list[float]]) -> list[float]:
    """Step 5: each number of the output is the weights times that number of every value, added up."""
    return [sum(w * v[c] for w, v in zip(weights, values)) for c in range(len(values[0]))]


def attend(vectors: list[list[float]], head: dict, scale: bool = True) -> dict:
    """Run one head over a sentence. Returns the queries, keys, values, scores, weights and outputs, one row per word."""
    queries = [times(x, head["wq"]) for x in vectors]       # step 1: what each word looks for ...
    keys = [times(x, head["wk"]) for x in vectors]          # ... the label each word holds up ...
    values = [times(x, head["wv"]) for x in vectors]        # ... and the note it carries
    scores = [scores_for(queries[i], keys, i, scale) for i in range(len(vectors))]
    weights = [softmax(row) for row in scores]
    outputs = [weighted_average(w, values) for w in weights]
    return {"queries": queries, "keys": keys, "values": values, "scores": scores, "weights": weights, "outputs": outputs}


def multi_head(vectors: list[list[float]], heads: list[dict]) -> list[list[float]]:
    """Several heads side by side: each word's outputs from every head, joined into one longer list."""
    runs = [attend(vectors, head)["outputs"] for head in heads]
    return [sum((run[i] for run in runs), []) for i in range(len(vectors))]


# --- what the page shows ---
def sentence(head: str, i: int) -> str:
    """The one line under the bars: where word i puts its weight."""
    w = attend(VECTORS, HEADS[head])["weights"][i]
    token = TOKENS[i]
    looked = [j for j in range(i + 1) if w[j] > 0]
    if len(looked) > 1 and max(w[:i + 1]) - min(w[:i + 1]) < 1e-9:
        return f"“{token}” spreads its weight evenly: {100 * w[0]:.0f}% on each of {len(looked)} words."
    top = max(range(len(w)), key=lambda j: w[j])
    return f"“{token}” puts {100 * w[top]:.0f}% of its weight on “{TOKENS[top]}”."


def blend_line(head: str, i: int) -> str:
    """The line about what the blend carries: the first number of a word's output says how much "an animal" it holds."""
    out = attend(VECTORS, HEADS[head])["outputs"][i]
    return f"Blended note, “an animal”: {out[0]:.2f}. “{TOKENS[i]}” alone: {CONTENT[i][0]:.2f}."


def percent(w: float) -> str:
    """A weight as a whole percent; a weight that is tiny but not zero shows as <1 so it never reads as zero."""
    return "<1" if 0 < w < 0.005 else f"{100 * w:.0f}"


def view(head: str, i: int) -> dict:
    """Everything the page writes for one head and one selected word. A word after word i is not looked at, so its cells are empty."""
    run = attend(VECTORS, HEADS[head])
    return {
        "weights": [percent(w) if j <= i else "" for j, w in enumerate(run["weights"][i])],
        "scores": [f"{s:.2f}" if j <= i else "" for j, s in enumerate(run["scores"][i])],
        "sentence": sentence(head, i), "blend": blend_line(head, i),
        "grid": [[percent(w) if j <= r else "" for j, w in enumerate(row)] for r, row in enumerate(run["weights"])],
    }


# --- why the division by sqrt(d) ---
def scaling_experiment(lengths=(4, 64, 1024), words=8, trials=100, seed=6) -> list[dict]:
    """Random queries and keys with numbers of mean 0 and spread 1, the assumption in the footnote in section 3.2.1 of the paper. For each
    length d, score one query against `words` keys, with and without the division, and report how spread out the scores are (their
    standard deviation) and how much weight the biggest score gets (the average over trials)."""
    rng = random.Random(seed)
    rows = []
    for d in lengths:
        raw, scaled, top_raw, top_scaled = [], [], [], []
        for _ in range(trials):
            q = [rng.gauss(0, 1) for _ in range(d)]
            s = [dot(q, [rng.gauss(0, 1) for _ in range(d)]) for _ in range(words)]
            raw += s
            scaled += [x / math.sqrt(d) for x in s]
            top_raw.append(max(softmax(s)))
            top_scaled.append(max(softmax([x / math.sqrt(d) for x in s])))
        rows.append({"d": d, "spread_raw": spread(raw), "spread_scaled": spread(scaled),
                     "top_raw": sum(top_raw) / trials, "top_scaled": sum(top_scaled) / trials})
    return rows


def spread(numbers: list[float]) -> float:
    mean = sum(numbers) / len(numbers)
    return math.sqrt(sum((x - mean) ** 2 for x in numbers) / len(numbers))


# --- what attention costs ---
LAYERS = 32     # copied from lesson 23 (prefill_sim.py): Llama 3.1 8B has 32 layers and 32 query heads per layer (they share 8 key-value heads)
HEADS_PER_LAYER = 32


def score_count(n: int) -> int:
    """Scores one head makes over a prompt of n words: word 1 scores 1 word, word 2 scores 2, ... word n scores n."""
    return n * (n + 1) // 2


def config() -> dict:
    """What the page needs to run the same steps: the words, the numbers per word, the maps of each head and the labels."""
    return {"tokens": TOKENS, "about": ABOUT, "vectors": VECTORS, "heads": HEADS, "labels": HEAD_LABELS,
            "opening": {"head": "refers", "token": TOKENS.index("it")}}