"""Reproduce Article 06 v2's constructed examples and local gradient checks.

Run: python assets/code/first-principles-06/verify.py
Dependency: numpy
These checks verify arithmetic and algebraic identities for teaching examples.
They are not trained-model experiments, linguistic evidence, or a site-render test.
"""
from __future__ import annotations

from collections.abc import Callable
import json
import numpy as np
from numpy.typing import NDArray

Array = NDArray[np.float64]


def softmax_rows(scores: Array) -> Array:
    """Stable last-axis Softmax; reject rows with no finite candidate."""
    scores = np.asarray(scores, dtype=np.float64)
    if scores.ndim < 1 or scores.shape[-1] == 0:
        raise ValueError("At least one candidate is required.")
    if np.isnan(scores).any() or np.isposinf(scores).any():
        raise ValueError("Scores cannot contain NaN or positive infinity.")
    if not np.isfinite(scores).any(axis=-1).all():
        raise ValueError("Every row needs at least one finite candidate.")
    shifted = scores - scores.max(axis=-1, keepdims=True)
    weights = np.exp(shifted)
    return weights / weights.sum(axis=-1, keepdims=True)


def attention(q: Array, k: Array, v: Array,
              visible: NDArray[np.bool_] | None = None) -> tuple[Array, Array]:
    """One attention read, with True meaning a visible key position."""
    if any(a.ndim != 2 for a in (q, k, v)):
        raise ValueError("Q, K, and V must be matrices.")
    if q.shape[1] == 0 or q.shape[1] != k.shape[1] or k.shape[0] != v.shape[0]:
        raise ValueError("Incompatible Query, Key, or Value shapes.")
    scores = q @ k.T / np.sqrt(q.shape[1])
    if visible is not None:
        if visible.shape != scores.shape:
            raise ValueError("Mask shape must equal the score-matrix shape.")
        scores = np.where(visible, scores, -np.inf)
    weights = softmax_rows(scores)
    return weights @ v, weights


def finite_difference(fn: Callable[[Array], float], x: Array,
                      step: float = 1e-6) -> Array:
    result = np.empty_like(x)
    for index in np.ndindex(x.shape):
        plus, minus = x.copy(), x.copy()
        plus[index] += step
        minus[index] -= step
        result[index] = (fn(plus) - fn(minus)) / (2.0 * step)
    return result


def main() -> None:
    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.])
    q, k, v = x @ wq, x @ wk, x @ wv
    out, weights = attention(q, k, v)
    visible = np.tril(np.ones((3, 3), dtype=bool))
    causal_out, causal_weights = attention(q, k, v, visible)
    checks: list[str] = []

    for source, projection, target in ((x, wq, q), (x, wk, k), (x, wv, v)):
        np.testing.assert_allclose(
            np.stack([projection.T @ row for row in source]), target
        )
    checks.append("Column-vector projections equal batched row-matrix projections")

    np.testing.assert_array_equal(q @ k.T, [[1, 0, 1], [1, 1, 2], [2, 1, 3]])
    np.testing.assert_allclose(weights[1], [.248255, .248255, .503490], atol=1e-6)
    np.testing.assert_allclose(out[1], [1.503490, .751745], atol=1e-6)
    np.testing.assert_allclose(causal_weights[1], [.5, .5, 0.], atol=1e-12)
    np.testing.assert_allclose(causal_out[1], [1., .5], atol=1e-12)
    np.testing.assert_allclose(weights.sum(axis=1), 1.)
    checks.append("Three-position numerical example, with and without causal masking")

    scalar_scores = np.array([2., 1.])
    scalar_values = np.array([10., 20.])
    np.testing.assert_allclose(softmax_rows(scalar_scores), softmax_rows(scalar_scores + 100.))
    assert scalar_scores @ scalar_values == 40.
    assert (scalar_scores + 100.) @ scalar_values == 3040.
    np.testing.assert_allclose(softmax_rows(scalar_scores) @ scalar_values, 12.6894142137)
    checks.append("Raw mixing versus shift-invariant normalized reading")

    for i in range(3):
        for j in range(3):
            for m in range(3):
                np.testing.assert_allclose(
                    np.log(weights[i, j] / weights[i, m]),
                    q[i] @ (k[j] - k[m]) / np.sqrt(2.), atol=1e-12
                )
    checks.append("Log weight ratio equals scaled Query dot Key-difference")

    symmetric_scores = np.array([[1., 2.], [2., 4.]])
    asymmetric_weights = softmax_rows(symmetric_scores)
    assert np.allclose(symmetric_scores, symmetric_scores.T)
    assert not np.allclose(asymmetric_weights, asymmetric_weights.T)
    checks.append("Symmetric scores need not produce symmetric normalized weights")

    masked_after = np.where(visible, weights, 0.)
    renormalized = masked_after / masked_after.sum(axis=1, keepdims=True)
    np.testing.assert_allclose(renormalized, causal_weights, atol=1e-12)
    assert not np.allclose(masked_after[1], causal_weights[1])
    try:
        softmax_rows(np.array([[-np.inf, -np.inf]]))
    except ValueError:
        pass
    else:
        raise AssertionError("An all-masked row should be rejected.")
    checks.append("Masking equivalence requires renormalization and nonempty visibility")

    order = np.array([2, 0, 1])
    permuted_out, _ = attention(q[order], k[order], v[order])
    np.testing.assert_allclose(permuted_out, out[order], atol=1e-12)
    checks.append("Permutation equivariance without a position mechanism or mask")

    changed_x = x.copy()
    changed_x[0] = [3., -.5]
    changed_q, changed_k, changed_v = changed_x @ wq, changed_x @ wk, changed_x @ wv
    np.testing.assert_allclose(changed_q[1], q[1])
    changed_out, _ = attention(changed_q, changed_k, changed_v)
    assert not np.allclose(changed_out[1], out[1])
    different_values_out, same_weights = attention(q, k, v + 1.)
    np.testing.assert_allclose(same_weights, weights)
    assert not np.allclose(different_values_out, out)
    checks.append("An unchanged Query can read different context; Values do not set weights")

    upstream = np.array([[.3, -.7], [1.2, .4], [-.2, .9]])
    max_gradient_error = 0.
    for mask in (None, visible):
        current_out, a = attention(q, k, v, mask)
        da = upstream @ v.T
        de = a * (da - np.sum(a * da, axis=1, keepdims=True))
        local_de = a * ((upstream @ v.T) - np.sum(upstream * current_out, axis=1, keepdims=True))
        np.testing.assert_allclose(de, local_de, atol=1e-12)
        dq = de @ k / np.sqrt(2.)
        dk = de.T @ q / np.sqrt(2.)
        dv = a.T @ upstream
        fns = (
            lambda z: float(np.sum(attention(z, k, v, mask)[0] * upstream)),
            lambda z: float(np.sum(attention(q, z, v, mask)[0] * upstream)),
            lambda z: float(np.sum(attention(q, k, z, mask)[0] * upstream)),
        )
        for value, analytic, fn in zip((q, k, v), (dq, dk, dv), fns):
            numeric = finite_difference(fn, value)
            np.testing.assert_allclose(analytic, numeric, rtol=2e-6, atol=2e-8)
            max_gradient_error = max(max_gradient_error, float(np.max(np.abs(analytic - numeric))))
        w_fns = (
            lambda z: float(np.sum(attention(x @ z, k, v, mask)[0] * upstream)),
            lambda z: float(np.sum(attention(q, x @ z, v, mask)[0] * upstream)),
            lambda z: float(np.sum(attention(q, k, x @ z, mask)[0] * upstream)),
        )
        for value, analytic, fn in zip((wq, wk, wv), (x.T @ dq, x.T @ dk, x.T @ dv), w_fns):
            numeric = finite_difference(fn, value)
            np.testing.assert_allclose(analytic, numeric, rtol=2e-6, atol=2e-8)
            max_gradient_error = max(max_gradient_error, float(np.max(np.abs(analytic - numeric))))
    checks.append("Local score gradients and Q/K/V and shared-projection gradients, by finite differences")
    print(json.dumps({
        "status": "PASS", "check_groups": checks,
        "row2_weights": weights[1].tolist(), "row2_output": out[1].tolist(),
        "causal_row2_weights": causal_weights[1].tolist(),
        "causal_row2_output": causal_out[1].tolist(),
        "max_gradient_absolute_error": max_gradient_error,
        "scope": "Constructed numerical checks only; no training or website-render validation."
    }, indent=2))


if __name__ == "__main__":
    main()