"""Checked teaching examples for Article 05: attention before QKV.

The states and parameters below are hand-specified. This script does not train
an encoder, decoder, or translation model, and the graph is not a training trace.
"""
from __future__ import annotations

import argparse
import json
from pathlib import Path
from xml.etree import ElementTree

import matplotlib.pyplot as plt
import numpy as np
from numpy.typing import NDArray

Array = NDArray[np.float64]
ROOT = Path(__file__).resolve().parent
H: Array = np.array([[2.0, 0.0], [1.0, -1.0], [0.0, -2.0]], dtype=np.float64)
W_S: Array = np.ones((2, 1), dtype=np.float64)
W_H: Array = np.eye(2, dtype=np.float64)
V: Array = np.array([1.0, -1.0], dtype=np.float64)
B: Array = np.zeros(2, dtype=np.float64)


def softmax(scores: Array) -> Array:
    """Stable softmax over a nonempty vector of finite scores."""
    if scores.ndim != 1 or not scores.size or not np.isfinite(scores).all():
        raise ValueError("Expected a nonempty, finite score vector.")
    weights = np.exp(scores - np.max(scores))
    return weights / weights.sum()


def attention(state: float, memory: Array = H) -> tuple[Array, Array, Array]:
    """Additive scores, source-position weights, and the weighted read.

    Math uses column vectors. Rows in `memory` simply store the h_j vectors.
    """
    if memory.ndim != 2 or memory.shape[1] != 2 or memory.shape[0] == 0:
        raise ValueError("Expected nonempty memory with shape [T, 2].")
    request = W_S @ np.array([state], dtype=np.float64)
    scores = np.tanh(memory @ W_H.T + request + B) @ V
    weights = softmax(scores)
    context = weights @ memory
    return scores, weights, context


def verify() -> dict[str, object]:
    examples: dict[str, object] = {}
    for state in (-1.0, 0.0, 1.0):
        scores, weights, context = attention(state)
        assert (weights > 0).all()
        np.testing.assert_allclose(weights.sum(), 1.0, atol=1e-14)
        np.testing.assert_allclose(softmax(scores + 13.0), weights, atol=1e-14)
        examples[str(state)] = {
            "scores": scores.tolist(),
            "weights": weights.tolist(),
            "context": context.tolist(),
        }

    np.testing.assert_allclose(attention(-1.0)[1], attention(1.0)[1][::-1], atol=1e-14)
    np.testing.assert_allclose(attention(0.0)[2], [1.0, -1.0], atol=1e-14)

    # Verify the partial derivative through a weighted read, holding H fixed.
    scores, weights, context = attention(-1.0)
    upstream = np.array([0.4, -0.7], dtype=np.float64)
    analytical = weights * ((H - context) @ upstream)
    eps = 1e-6
    numerical = np.zeros_like(scores)
    for j in range(scores.size):
        delta = np.zeros_like(scores)
        delta[j] = eps
        loss_plus = float(upstream @ (softmax(scores + delta) @ H))
        loss_minus = float(upstream @ (softmax(scores - delta) @ H))
        numerical[j] = (loss_plus - loss_minus) / (2.0 * eps)
    np.testing.assert_allclose(analytical, numerical, atol=2e-10, rtol=1e-7)
    np.testing.assert_allclose(analytical.sum(), 0.0, atol=1e-14)

    # Remove tanh: the state-dependent term is a shared offset and cancels.
    def linear_weights(state: float) -> Array:
        return softmax((H @ W_H.T + W_S[:, 0] * state + B) @ V)
    np.testing.assert_allclose(linear_weights(-2.0), linear_weights(1.7), atol=1e-14)

    # If every item is identical, the read is independent of normalized weights.
    identical = np.repeat(np.array([[0.2, -0.3]]), repeats=3, axis=0)
    np.testing.assert_allclose(attention(-1.0, identical)[2], identical[0], atol=1e-14)

    # Independent double-precision autodiff check of the full additive module.
    import torch
    tensors = tuple(
        torch.tensor(a, dtype=torch.float64, requires_grad=True)
        for a in (np.array([-1.0]), H, W_S, W_H, V, B)
    )

    def torch_read(
        state: torch.Tensor,
        memory: torch.Tensor,
        w_s: torch.Tensor,
        w_h: torch.Tensor,
        v: torch.Tensor,
        bias: torch.Tensor,
    ) -> torch.Tensor:
        logits = torch.tanh(memory @ w_h.T + w_s @ state + bias) @ v
        return logits.softmax(dim=0) @ memory

    passed = torch.autograd.gradcheck(torch_read, tensors, eps=1e-6, atol=1e-6, rtol=1e-4)
    assert passed
    np.testing.assert_allclose(torch_read(*tensors).detach().numpy(), context, atol=1e-14)

    return {
        "description": "Constructed additive-attention examples; no translation training.",
        "memory": H.tolist(),
        "W_s": W_S.tolist(),
        "W_h": W_H.tolist(),
        "v": V.tolist(),
        "bias": B.tolist(),
        "states": examples,
        "score_gradient": {
            "upstream": upstream.tolist(),
            "analytical": analytical.tolist(),
            "finite_difference": numerical.tolist(),
            "max_abs_error": float(np.abs(analytical - numerical).max()),
        },
        "torch_gradcheck": bool(passed),
    }


def plot() -> None:
    """Plot a state sweep, not a sequence of optimization steps."""
    states = np.linspace(-3.0, 3.0, 601)
    weights = np.array([attention(float(s))[1] for s in states])
    fig, ax = plt.subplots(figsize=(9.0, 5.2))
    styles = ("-", "--", "-.")
    labels = ("Position 1: I", "Position 2: like", "Position 3: cats")
    for j, (style, label) in enumerate(zip(styles, labels)):
        ax.plot(states, weights[:, j], linestyle=style, label=label, linewidth=2.3)
    ax.axvline(-1.0, linestyle=":", linewidth=1.0, alpha=0.4)
    ax.axvline(1.0, linestyle=":", linewidth=1.0, alpha=0.4)
    ax.set(
        xlabel="Read-request state s (a constructed scalar)",
        ylabel="Source-position attention weight",
        title="Same memory and parameters, different reading weights",
        xlim=(-3.0, 3.0),
        ylim=(0.0, 0.65),
    )
    ax.grid(alpha=0.18)
    ax.legend(loc="upper center", bbox_to_anchor=(0.5, -0.19), ncol=3, frameon=False)
    fig.subplots_adjust(left=0.11, right=0.98, top=0.9, bottom=0.26)
    out = ROOT / "figures"
    out.mkdir(exist_ok=True)
    fig.savefig(out / "01-state-dependent-reading.svg", bbox_inches="tight")
    fig.savefig(out / "01-state-dependent-reading.png", dpi=180, bbox_inches="tight")
    ElementTree.parse(out / "01-state-dependent-reading.svg")
    plt.show()
    plt.close(fig)


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--verify-only", action="store_true", help="Check calculations without writing or displaying figures.")
    args = parser.parse_args()
    result = verify()
    if args.verify_only:
        print(json.dumps(result, indent=2))
        return
    (ROOT / "verification.json").write_text(json.dumps(result, indent=2), encoding="utf-8")
    plot()
    print("Scores, normalization, context vectors, shift invariance and gradients verified.")
    print("Full additive module: PyTorch gradcheck passed. SVG: XML parsing passed.")


if __name__ == "__main__":
    main()