"""How big the KV cache is, and what it saves.

Standard library plus numpy. Run `python3 show_numbers.py` to print every
number the lesson quotes.
"""
from dataclasses import dataclass

import numpy as np


@dataclass(frozen=True)
class ModelConfig:
    name: str
    n_layers: int
    n_kv_heads: int      # key/value heads, not query heads
    head_dim: int
    bytes_per_elem: int  # 2 for 16-bit floats (fp16 or bf16)
    max_tokens: int      # longest sequence the model accepts


# Llama 3 8B: 32 layers, 8 key/value heads (grouped-query attention),
# hidden size 4096 / 32 query heads = head_dim 128, bf16, 8,192-token context.
LLAMA3_8B = ModelConfig("Llama 3 8B", 32, 8, 128, 2, 8192)
# OPT-13B: 40 layers, 40 heads of size 128 (hidden size 5120), fp16, 2,048 tokens.
OPT_13B = ModelConfig("OPT-13B", 40, 40, 128, 2, 2048)


def kv_bytes_per_token(cfg: ModelConfig) -> int:
    """2 (one key, one value) x layers x kv_heads x head_dim x bytes."""
    return 2 * cfg.n_layers * cfg.n_kv_heads * cfg.head_dim * cfg.bytes_per_elem


def kv_bytes(cfg: ModelConfig, tokens: int) -> int:
    return kv_bytes_per_token(cfg) * tokens


def pairs_without_cache(n_steps: int) -> int:
    """Step t must have keys and values for all t positions: 1 + 2 + ... + n."""
    return n_steps * (n_steps + 1) // 2


def pairs_with_cache(n_steps: int) -> int:
    """Each step computes the key and value of the one new position."""
    return n_steps


def _softmax(s):
    e = np.exp(s - s.max())
    return e / e.sum()


def decode_both_ways(n_steps: int, d: int = 8, seed: int = 0):
    """One attention head generating n_steps positions, with and without a cache.

    Returns (outputs_without, outputs_with, pairs_without, pairs_with).
    A "pair" is one key vector plus one value vector computed for one position.
    The input vectors are fixed random numbers; a real model feeds each output
    back in, but the counting is the same.
    """
    rng = np.random.default_rng(seed)
    w_q, w_k, w_v = rng.standard_normal((3, d, d)) / np.sqrt(d)
    x = rng.standard_normal((n_steps, d))

    out_a, pairs_a = [], 0
    for t in range(1, n_steps + 1):
        keys, values = x[:t] @ w_k, x[:t] @ w_v       # recompute the whole past
        pairs_a += t
        q = x[t - 1] @ w_q
        out_a.append(_softmax(keys @ q / np.sqrt(d)) @ values)

    out_b, pairs_b = [], 0
    k_cache, v_cache = [], []
    for t in range(1, n_steps + 1):
        k_cache.append(x[t - 1] @ w_k)                 # only the new position
        v_cache.append(x[t - 1] @ w_v)
        pairs_b += 1
        q = x[t - 1] @ w_q
        keys, values = np.array(k_cache), np.array(v_cache)
        out_b.append(_softmax(keys @ q / np.sqrt(d)) @ values)
    return np.array(out_a), np.array(out_b), pairs_a, pairs_b