"""Talent to Cash: original one-dimensional attention example. Standard library.
Run: python causal-attention.py
Q=K=V and d_k=1. This is a concept check, not a trained Transformer.
"""
from math import exp, isclose


def attention(values, causal):
    weights, outputs = [], []
    for row, query in enumerate(values):
        scores = [query * key if not causal or col <= row else -float("inf")
                  for col, key in enumerate(values)]
        maximum = max(scores)
        exponentials = [exp(score - maximum) for score in scores]
        total = sum(exponentials)
        line = [value / total for value in exponentials]
        weights.append(line)
        outputs.append(sum(weight * value for weight, value in zip(line, values)))
    return weights, outputs


if __name__ == "__main__":
    for causal in [True, False]:
        weights, before = attention([1, 2, 3], causal)
        _, after = attention([1, 2, 5], causal)
        assert all(isclose(sum(row), 1.0) for row in weights)
        if causal:
            assert all(isclose(before[i], after[i]) for i in [0, 1])
        else:
            assert any(not isclose(before[i], after[i]) for i in [0, 1])
        print("causal:", causal, "before:", before, "after:", after)