"""One constructed example shared by Article 06 figures and ManimGX scenes."""
import numpy as np

X = np.array([[1., 0.], [0., 1.], [1., 1.]])
WQ = np.eye(2)
WK = np.array([[1., 1.], [0., 1.]])
WV = np.diag([2., 1.])


def softmax(scores):
    scores = np.asarray(scores, dtype=float)
    if not np.isfinite(scores).any(axis=-1).all():
        raise ValueError("Each row needs a visible candidate.")
    exp = np.exp(scores - scores.max(axis=-1, keepdims=True))
    return exp / exp.sum(axis=-1, keepdims=True)


def checkpoint():
    q, k, v = X @ WQ, X @ WK, X @ WV
    raw = q @ k.T
    scores = raw / np.sqrt(q.shape[1])
    weights = softmax(scores)
    visible = np.tril(np.ones(raw.shape, dtype=bool))
    masked = np.where(visible, scores, -np.inf)
    causal = softmax(masked)
    return dict(x=X, wq=WQ, wk=WK, wv=WV, q=q, k=k, v=v, raw=raw,
                scores=scores, weights=weights, out=weights @ v,
                masked=masked, causal=causal, causal_out=causal @ v)