#!/usr/bin/env python3
"""The attention plate's NumPy witness (Python 3 + NumPy).

python3 attention_specimen.py --query 2 --tau 1
python3 attention_specimen.py --fixture > attentionSpecimen.reference.json

Query indices are zero-based. Fixture scores use JSON null for masked -infinity.
The small vectors are constructed examples, not learned representations.
"""

import argparse
import hashlib
import json
import math
from pathlib import Path

import numpy as np

TOKENS = ["the", "cat", "sat", "on"]
Q = np.array([[1, 0], [-1, 0.5], [1.5, 0.5], [-1, 1.5]], dtype=float)
K = np.array([[1, 0], [2, 1], [-1, 1], [0, -1]], dtype=float)
V = np.array([[-1, 1], [1, 2], [2, -1], [-1, -1]], dtype=float)
TEMPERATURES = [0.25, 0.5, 1, 2, 4]


def display_number(value, digits=3):
    rounded = f"{value:.{digits}f}"
    return f"{0:.{digits}f}" if float(rounded) == 0 else rounded


def compute(query_index=2, tau=1):
    if type(query_index) is not int or not 0 <= query_index < len(TOKENS):
        raise ValueError("Query index must be an integer from 0 to 3.")
    if not math.isfinite(tau) or tau <= 0:
        raise ValueError("Temperature must be finite and positive.")
    scores = K @ Q[query_index] / np.sqrt(K.shape[1])
    scores[query_index + 1:] = -np.inf
    # The negative overflow at extremely small tau represents exp(-infinity) = 0.
    with np.errstate(over="ignore"):
        masses = np.exp((scores - scores.max()) / tau)
    weights = masses / masses.sum()
    output = weights @ V
    return {
        "queryIndex": query_index,
        "tau": tau,
        "scores": [None if np.isneginf(s) else float(s) for s in scores],
        "weights": weights.tolist(),
        "output": output.tolist(),
        "display": {
            "scores": ["−∞" if np.isneginf(s) else display_number(s) for s in scores],
            "weights": [display_number(a * 100, 1) + "%" for a in weights],
            "output": [display_number(v) for v in output],
        },
    }


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--query", type=int, default=2)
    parser.add_argument("--tau", type=float, default=1)
    parser.add_argument("--fixture", action="store_true")
    args = parser.parse_args()
    if args.fixture:
        result = {
            "generator": "public/specimen/attention_specimen.py --fixture (NumPy)",
            "witnessSha256": hashlib.sha256(Path(__file__).read_bytes()).hexdigest(),
            "tokens": TOKENS,
            "queries": Q.tolist(), "keys": K.tolist(), "values": V.tolist(),
            "cases": [compute(i, tau) for i in range(len(TOKENS)) for tau in TEMPERATURES],
        }
        print(json.dumps(result, indent=2, ensure_ascii=False, allow_nan=False))
    else:
        result = compute(args.query, args.tau)
        print(f"Query {TOKENS[args.query]} (position {args.query + 1}), temperature {args.tau:g}")
        for token, score, weight in zip(TOKENS, result["display"]["scores"], result["display"]["weights"]):
            print(f"{token:>3}: score {score:>6}; weight {weight:>6}")
        print("o = (" + ", ".join(result["display"]["output"]) + ")")


if __name__ == "__main__":
    main()