"""Online softmax and one-pass attention for a single query, in plain NumPy.

Run `python3 online_softmax.py` to print every number the lesson quotes.
"""
import math

import numpy as np


class SlowRow:
    """A row of scores that lives in "slow memory".

    Looping over it hands out `block` scores at a time. Each score handed out
    counts as one read, so one full loop (one pass) adds len(scores) to `reads`.
    """

    def __init__(self, scores, block=1):
        self.scores = np.asarray(scores)
        self.block = block
        self.reads = 0

    def __iter__(self):
        for i in range(0, len(self.scores), self.block):
            chunk = self.scores[i:i + self.block]
            self.reads += len(chunk)
            yield chunk


# --- softmax of one row ----------------------------------------------------

def naive_softmax(x):
    """Textbook formula. Overflows as soon as one exp() does."""
    e = np.exp(x)
    return e / e.sum()


def stable_softmax(row):
    """Subtract the max first. Needs three passes over the row."""
    m = max(chunk.max() for chunk in row)                         # pass 1: max
    total = sum(np.exp(chunk - m).sum() for chunk in row)         # pass 2: sum
    return np.concatenate([np.exp(c - m) / total for c in row])   # pass 3: output


def online_stats(row):
    """Return (m, l) in ONE pass: m = max(x), l = sum(exp(x - m))."""
    m, l = -np.inf, 0.0
    for chunk in row:
        m_new = max(m, chunk.max())
        if m_new == -np.inf:          # a block that is entirely -inf adds nothing
            continue
        l = l * np.exp(m - m_new) + np.exp(chunk - m_new).sum()
        m = m_new
    return m, l


def online_softmax(row):
    """One pass for (m, l), one more pass to write the probabilities."""
    m, l = online_stats(row)
    return np.concatenate([np.exp(c - m) / l for c in row])


# --- attention for one query ----------------------------------------------

def attention_reference(q, K, V):
    """softmax(K q / sqrt(d)) @ V, computed the stable way in float64."""
    s = K.astype(np.float64) @ q.astype(np.float64) / math.sqrt(len(q))
    p = np.exp(s - s.max())
    return (p / p.sum()) @ V.astype(np.float64)


def attention_state(q, K, V, block):
    """Stream over K and V `block` rows at a time. Return the state (m, l, acc).

    Only one block of scores exists at any moment; the full score row never does.
    """
    m, l = -np.inf, 0.0
    acc = np.zeros(V.shape[1], dtype=V.dtype)
    for i in range(0, len(K), block):
        s = K[i:i + block] @ q / math.sqrt(len(q))   # scores for this block only
        m_new = max(m, s.max())
        if m_new == -np.inf:                          # fully masked so far: adds nothing
            continue
        alpha = np.exp(m - m_new)                     # 1.0 when the max did not move
        p = np.exp(s - m_new)
        l = l * alpha + p.sum()
        acc = acc * alpha + p @ V[i:i + block]        # same alpha as the sum
        m = m_new
    return m, l, acc


def online_attention(q, K, V, block):
    """softmax(K q / sqrt(d)) @ V in one pass over K and V."""
    _, l, acc = attention_state(q, K, V, block)
    return acc / l


def merge(a, b):
    """Combine two states built from different parts of the row."""
    (m_a, l_a, acc_a), (m_b, l_b, acc_b) = a, b
    m = max(m_a, m_b)
    if m == -np.inf:
        return a
    w_a, w_b = np.exp(m_a - m), np.exp(m_b - m)
    return m, l_a * w_a + l_b * w_b, acc_a * w_a + acc_b * w_b


# --- numbers for the lesson ------------------------------------------------

def score_matrix_bytes(n, bytes_per_score=4):
    """Size of an n x n score matrix for one head, in bytes."""
    return n * n * bytes_per_score


WORKED_X = [1.0, 3.0, 2.0, 5.0, 4.0]
WORKED_V = [10.0, 20.0, 30.0, 40.0, 50.0]


def trace(x, v, by_hand=False):
    """One row per element: the state after reading that element (block size 1).

    With `by_hand=True` the arithmetic matches the lesson's tables: alpha and p are
    exp() rounded to 5 places, and l and acc are rounded to 4 places after every row,
    so each row follows from the row above using only the digits shown.
    """
    def rnd(t, places):
        return round(t, places) if by_hand else t

    m, l, acc, rows = -math.inf, 0.0, 0.0, []
    for i, (xi, vi) in enumerate(zip(x, v), start=1):
        m_new = max(m, xi)
        alpha = rnd(math.exp(m - m_new), 5)
        p = rnd(math.exp(xi - m_new), 5)
        l = rnd(l * alpha + p, 4)
        acc = rnd(acc * alpha + p * vi, 4)
        m = m_new
        rows.append(dict(step=i, x=xi, v=vi, m=m, alpha=alpha, p=p, l=l, acc=acc,
                         out=acc / l))
    return rows


def _show():
    f32 = np.float32
    big = np.array([1000.0, 1001.0, 1002.0], dtype=f32)
    with np.errstate(all="ignore"):
        print("float32 exp(88.7) =", np.exp(f32(88.7)), "  exp(88.8) =", np.exp(f32(88.8)))
        print("float32 exp(1000) =", np.exp(f32(1000)))
        print("float64 exp(709) =", np.exp(709.0), "  exp(710) =", np.exp(710.0))
        print("naive_softmax([1000,1001,1002]) =", naive_softmax(big))
    print("stable_softmax([1000,1001,1002]) =", stable_softmax(SlowRow(big)))
    with np.errstate(all="ignore"):
        print("float16 exp(11.0) =", np.exp(np.float16(11.0)), "  exp(11.1) =", np.exp(np.float16(11.1)))
    print("float16 max =", np.finfo(np.float16).max, " float32 max =", np.finfo(np.float32).max)
    print("4096 x 4096 float32 scores =", score_matrix_bytes(4096) / 2**20, "MiB")
    print("32768 x 32768 float32 scores =", score_matrix_bytes(32768) / 2**30, "GiB")

    print("\nworked example, x =", WORKED_X, " v =", WORKED_V)
    for r in trace(WORKED_X, WORKED_V):
        print("  step %d x=%g m=%g alpha=%.4f p=%.4f l=%.4f acc=%.4f out=%.4f" % (
            r["step"], r["x"], r["m"], r["alpha"], r["p"], r["l"], r["acc"], r["out"]))
    final = trace(WORKED_X, WORKED_V)[-1]
    probs = np.exp(np.array(WORKED_X) - final["m"]) / final["l"]
    print("  final probabilities:", np.round(probs, 4), " sum =", probs.sum())
    print("  naive weighted average:", probs @ np.array(WORKED_V))

    print("\nsame example, by hand (5-place constants, rows rounded to 4 places)")
    hand = trace(WORKED_X, WORKED_V, by_hand=True)
    for r in hand:
        print("  step %d alpha=%.5f p=%.5f l=%.4f acc=%.4f out=%.2f" % (
            r["step"], r["alpha"], r["p"], r["l"], r["acc"], r["out"]))
    terms = [round(math.exp(xi - hand[-1]["m"]), 5) for xi in WORKED_X]
    print("  exp(x - 5):", terms, " sum =", round(sum(terms), 4))
    hand_probs = [round(t / hand[-1]["l"], 4) for t in terms]
    print("  probabilities:", hand_probs, " sum =", round(sum(hand_probs), 4))
    print("  probabilities . v =", round(sum(pr * vi for pr, vi in zip(hand_probs, WORKED_V)), 2))

    rng = np.random.default_rng(0)
    x = rng.standard_normal(1000) * 5
    err = np.abs(online_softmax(SlowRow(x, block=64)) - stable_softmax(SlowRow(x))).max()
    print("max |online - stable| over 1000 float64 scores:", err)
    q, K, V = rng.standard_normal(16), rng.standard_normal((1000, 16)), rng.standard_normal((1000, 16))
    err = np.abs(online_attention(q, K, V, 64) - attention_reference(q, K, V)).max()
    print("max |online attention - reference|, n=1000, d=16, block=64:", err)
    x = x.astype(f32)
    for name, fn in (("stable (3 passes)", stable_softmax),
                     ("online stats only", online_stats),
                     ("online softmax", online_softmax)):
        row = SlowRow(x, block=64)
        fn(row)
        print("reads for n=1000, %-18s %d  (= %.0f x n)" % (name, row.reads, row.reads / 1000))


if __name__ == "__main__":
    _show()