# examples/violin_attention.py
"""Standalone torch-only educational VIOLIN attention core.

Seeded Q/K/V and a synthetic backward pass, not a pretrained classifier,
VTAB accuracy reproduction, trained model, or speedup benchmark. Dense N*N
attention remains quadratic; eight curve masks share one QK product.

Reference: arXiv:2606.14757v1, equations 13-15, Appendices B.3/D.1/D.2/G.3.
Official code consulted: LIONS-EPFL/VIOLIN at
893f89c635a01d92529d6013450975862c956a7c; this is an original attention core.
The paper calls Morton bit-interleaving "Peano"; this example follows G.3 and
its official implementation (column bits even, row bits odd), matching our
JS demonstration. Transposition runs the curve on swapped dimensions first.

Beta is initialized uniformly in [5, 9]; alpha=1 is a teaching initialization.
Official pretraining initializes alpha from Normal(0, .5). The paper's printed
sigmoid range in D.1 is incorrect; sigmoid(5) is approximately .993307.
Geometry alone is cached; learned masks are rebuilt with autograd every call.
No data, network, projections, dropout, or additional special tokens are used.

Generalized Gilbert traversal reused from examples/hilbert_attention.py,
adapted from the authors' MIT-licensed hwt.py (Copyright (c) 2025 Yunge Li):
https://github.com/Yunge6666/Hilbert-Local-Attention/blob/main/hwt.py
The generalized Gilbert algorithm is cited as Cerveny (2024) by VIOLIN.

MIT License
Copyright (c) 2025 Yunge Li

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
"""

import argparse
import math
from functools import lru_cache

import torch
from torch import nn

CURVES = tuple(
    (curve, transposed)
    for curve in ("snake", "zigzag", "hilbert", "peano")
    for transposed in (False, True)
)


def _positive_integer(value, name):
    """Reject boolean/float dimensions, even when integer cache keys exist."""
    if isinstance(value, bool) or not isinstance(value, int) or value <= 0:
        raise ValueError(f"{name} must be a positive integer")


def _validate_configuration(height, width, heads):
    for value, name in ((height, "height"), (width, "width"), (heads, "heads")):
        _positive_integer(value, name)
    if height * width > 1024:
        raise ValueError(
            "height*width must be <=1024 patches (dense quadratic memory)"
        )
    if heads > 12:
        raise ValueError("heads must be <=12")


def _morton_key(flat_index, width):
    row, col = divmod(flat_index, width)
    key = 0
    for bit in range(max(row, col).bit_length()):
        key |= ((col >> bit) & 1) << (2 * bit)
        key |= ((row >> bit) & 1) << (2 * bit + 1)
    return key


def _generate2d(x, y, ax, ay, bx, by, result):
    """Generate the generalized Gilbert traversal, preserving signed halves."""
    width, height = abs(ax + ay), abs(bx + by)

    def sign(value):
        return (value > 0) - (value < 0)

    dax, day, dbx, dby = map(sign, (ax, ay, bx, by))
    if height == 1:
        result.extend((x + i * dax, y + i * day) for i in range(width))
        return
    if width == 1:
        result.extend((x + i * dbx, y + i * dby) for i in range(height))
        return
    ax2, ay2, bx2, by2 = ax // 2, ay // 2, bx // 2, by // 2
    width2, height2 = abs(ax2 + ay2), abs(bx2 + by2)
    if 2 * width > 3 * height:
        if width2 % 2 and width > 2:
            ax2, ay2 = ax2 + dax, ay2 + day
        _generate2d(x, y, ax2, ay2, bx, by, result)
        _generate2d(x + ax2, y + ay2, ax - ax2, ay - ay2, bx, by, result)
    else:
        if height2 % 2 and height > 2:
            bx2, by2 = bx2 + dbx, by2 + dby
        _generate2d(x, y, bx2, by2, ax2, ay2, result)
        _generate2d(x + bx2, y + by2, ax, ay, bx - bx2, by - by2, result)
        _generate2d(
            x + ax - dax + bx2 - dbx,
            y + ay - day + by2 - dby,
            -bx2,
            -by2,
            -(ax - ax2),
            -(ay - ay2),
            result,
        )


# region: curves
@lru_cache(maxsize=128, typed=True)
def curve_order(curve, height, width, transposed=False):
    """Return curve position -> row-major patch index for any rectangle."""
    _positive_integer(height, "height")
    _positive_integer(width, "width")
    if curve not in ("snake", "zigzag", "hilbert", "peano"):
        raise ValueError("curve must be snake, zigzag, hilbert, or peano")
    if not isinstance(transposed, bool):
        raise TypeError("transposed must be a boolean")
    if transposed:
        swapped = curve_order(curve, width, height)
        order = tuple(
            (index % height) * width + index // height for index in swapped
        )
    elif curve == "snake":
        order = tuple(
            row * width + col
            for row in range(height)
            for col in (
                range(width) if row % 2 == 0 else reversed(range(width))
            )
        )
    elif curve == "zigzag":
        coordinates = [
            (row, col) for row in range(height) for col in range(width)
        ]
        coordinates.sort(
            key=lambda rc: (sum(rc), -rc[0] if sum(rc) % 2 == 0 else rc[0])
        )
        order = tuple(row * width + col for row, col in coordinates)
    elif curve == "peano":  # The authors' name for Morton bit-interleaving.
        order = tuple(
            sorted(range(height * width), key=lambda i: _morton_key(i, width))
        )
    else:  # Generalized Gilbert supports non-power-of-two maps.
        path = []
        if width >= height:
            _generate2d(0, 0, width, 0, 0, height, path)
        else:
            _generate2d(0, 0, 0, height, width, 0, path)
        order = tuple(row * width + col for col, row in path)
    if sorted(order) != list(range(height * width)):
        raise ValueError("Curve traversal must be a permutation of every patch")
    return order


# endregion: curves


# region: ranks
@lru_cache(maxsize=128, typed=True)
def curve_ranks(curve, height, width, transposed=False):
    """Invert the order: row-major patch index -> its curve position."""
    order = curve_order(curve, height, width, transposed)
    ranks = [0] * len(order)
    for position, patch in enumerate(order):
        ranks[patch] = position
    return tuple(ranks)


# endregion: ranks


class ViolinAttention(nn.Module):
    """Fixed-map attention with only eight beta and one alpha per head.

    Move the module with .to(device/dtype) before providing Q/K/V. A module's
    map and class-token setting are fixed; create another module to resize.
    """

    # region: parameters
    def __init__(self, height, width, heads, class_token=False, seed=0):
        super().__init__()
        _validate_configuration(height, width, heads)
        if not isinstance(class_token, bool):
            raise TypeError("class_token must be a boolean")
        if isinstance(seed, bool) or not isinstance(seed, int):
            raise TypeError("seed must be an integer")
        self.height, self.width, self.heads = height, width, heads
        self.class_token = class_token
        generator = torch.Generator(device="cpu").manual_seed(seed)
        self.beta = nn.Parameter(
            torch.empty(heads, 8).uniform_(5, 9, generator=generator)
        )
        self.alpha = nn.Parameter(torch.ones(heads))  # Teaching initialization.
        ranks = [
            curve_ranks(curve, height, width, transposed)
            for curve, transposed in CURVES
        ]
        self.register_buffer("ranks", torch.tensor(ranks, dtype=torch.long))

    # endregion: parameters

    # region: decay
    def curve_masks(self):
        """Return differentiable [heads, eight curves, patches, patches]."""
        distances = (self.ranks[:, :, None] - self.ranks[:, None, :]).abs()
        gamma = torch.sigmoid(self.beta)
        # Ranks put each curve's KMS mask into the common row-major basis.
        return gamma[:, :, None, None] ** distances[None, :, :, :]

    # endregion: decay

    # region: average
    def patch_mask(self):
        """Average the eight aligned masks, before scores or softmax."""
        return self.curve_masks().mean(dim=1)

    # endregion: average

    # region: special-tokens
    def attention_mask(self):
        """Keep a leading class token's row/column equal to one (D.2)."""
        mask = self.patch_mask()
        if not self.class_token:
            return mask
        heads, patches, _ = mask.shape
        column = mask.new_ones(heads, patches, 1)
        row = mask.new_ones(heads, 1, patches + 1)
        return torch.cat((row, torch.cat((column, mask), dim=-1)), dim=-2)

    # endregion: special-tokens

    def _validate_qkv(self, q, k, v):
        expected_tokens = self.height * self.width + int(self.class_token)
        tensors = (q, k, v)
        if any(
            not isinstance(tensor, torch.Tensor) or tensor.ndim != 4
            for tensor in tensors
        ):
            raise ValueError(
                "Q/K/V must be [batch, heads, tokens, head_dim] tensors"
            )
        if any(tensor.shape != q.shape for tensor in tensors):
            raise ValueError("Q/K/V must have the same shape")
        if (
            any(dimension <= 0 for dimension in q.shape)
            or q.shape[1] != self.heads
            or q.shape[2] != expected_tokens
        ):
            raise ValueError(
                "Positive dimensions, configured heads, "
                "and matching token count required"
            )
        if any(
            not tensor.is_floating_point()
            or tensor.dtype != q.dtype
            or tensor.device != q.device
            for tensor in tensors
        ):
            raise ValueError(
                "Q/K/V must have the same floating dtype and device"
            )
        if self.beta.device != q.device or self.beta.dtype != q.dtype:
            raise ValueError(
                "Move the module to the Q/K/V dtype and device first"
            )
        if any(not torch.isfinite(tensor).all().item() for tensor in tensors):
            raise ValueError("Q/K/V values must be finite")

    # region: attention
    def forward(self, q, k, v):
        """One QK product; multiply logits by alpha*M before softmax."""
        self._validate_qkv(q, k, v)
        scores = (q @ k.transpose(-2, -1)) / math.sqrt(q.shape[-1])
        mask = self.attention_mask()
        # Negative scores shrink towards zero and can gain relative probability.
        weights = torch.softmax(
            scores * self.alpha[None, :, None, None] * mask, dim=-1
        )
        return weights @ v

    # endregion: attention


def permutation_reference(model):
    """Independent dense literal mask construction for the CLI comparison.

    Build the KMS mask in curve coordinates, then scatter both axes via the
    forward permutation. No inverse ranks or product distance tensor is used.
    This is equivalent to P.T @ base_mask @ P without cubic matrix products.
    """
    patches = model.height * model.width
    positions = torch.arange(patches, device=model.beta.device)
    base_distances = (positions[:, None] - positions[None, :]).abs()
    reference = model.beta.new_zeros(model.heads, patches, patches)
    for curve_index, (curve, transposed) in enumerate(CURVES):
        order = torch.tensor(
            curve_order(curve, model.height, model.width, transposed),
            device=model.beta.device,
        )
        for head in range(model.heads):
            gamma = 1 / (
                1 + math.exp(-float(model.beta[head, curve_index].detach()))
            )
            base = gamma ** base_distances.to(dtype=model.beta.dtype)
            reference[head, order[:, None], order[None, :]] += base / 8
    if model.class_token:
        extended = reference.new_ones(model.heads, patches + 1, patches + 1)
        extended[:, 1:, 1:] = reference
        reference = extended
    return reference


def main(argv=None):
    """Run a bounded float64 reference check and synthetic backward pass."""
    parser = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    parser.add_argument("--height", type=int, default=14)
    parser.add_argument("--width", type=int, default=14)
    parser.add_argument("--heads", type=int, default=3)
    parser.add_argument("--head-dim", type=int, default=16)
    parser.add_argument("--seed", type=int, default=7)
    parser.add_argument("--device", choices=("cpu", "cuda"), default="cpu")
    parser.add_argument("--class-token", action="store_true")
    args = parser.parse_args(argv)
    try:
        _validate_configuration(args.height, args.width, args.heads)
        _positive_integer(args.head_dim, "head-dim")
        if args.head_dim > 128:
            raise ValueError("head-dim must be <=128")
        if args.device == "cuda" and not torch.cuda.is_available():
            raise ValueError("CUDA requested but unavailable")
        model = ViolinAttention(
            args.height,
            args.width,
            args.heads,
            class_token=args.class_token,
            seed=args.seed,
        ).to(device=args.device, dtype=torch.float64)
    except (ValueError, RuntimeError) as error:
        parser.error(str(error))
    generator = torch.Generator(device="cpu").manual_seed(args.seed)
    tokens = args.height * args.width + int(args.class_token)
    shape = (1, args.heads, tokens, args.head_dim)
    q, k, v = (
        torch.randn(shape, generator=generator, dtype=torch.float64).to(
            args.device
        )
        for _ in range(3)
    )
    output = model(q, k, v)
    with torch.no_grad():
        mask = model.attention_mask()
        literal_mask = permutation_reference(model)
        scores = q @ k.transpose(-2, -1) / math.sqrt(args.head_dim)
        expected = (
            torch.softmax(
                scores * model.alpha[None, :, None, None] * literal_mask, -1
            )
            @ v
        )
        matches = torch.allclose(
            mask, literal_mask, atol=1e-12, rtol=1e-10
        ) and torch.allclose(output, expected, atol=1e-12, rtol=1e-10)
        symmetric = torch.allclose(mask, mask.transpose(-2, -1))
        diagonal = torch.equal(
            mask.diagonal(dim1=-2, dim2=-1), torch.ones_like(mask[:, :, 0])
        )
    target = torch.randn(shape, generator=generator, dtype=torch.float64).to(
        args.device
    )
    loss = (output - target).square().mean()
    loss.backward()
    gradient_checks = {}
    gradient_diagnostics = {}
    # One patch has only zero curve distances, so beta cannot affect its mask.
    # Alpha cannot affect a one-element softmax, but remains active with CLS.
    inactive_reasons = {
        "beta": (
            "all curve distances are zero"
            if args.height * args.width == 1
            else None
        ),
        "alpha": "one-element softmax is constant" if tokens == 1 else None,
    }
    for name in ("beta", "alpha"):
        gradient = getattr(model, name).grad
        finite = bool(torch.isfinite(gradient).all())
        nonzero = bool(gradient.abs().sum() > 0)
        inactive_reason = inactive_reasons[name]
        gradient_checks[name] = finite and (
            not nonzero if inactive_reason else nonzero
        )
        gradient_diagnostics[name] = (finite, nonzero, inactive_reason)
    print(
        "Educational dense attention core; synthetic tensors and loss, "
        "no accuracy benchmark."
    )
    print(
        "beta~Uniform[5,9]; alpha=1 teaching initialization; "
        "geometry fixed per module."
    )
    print(
        f"device={args.device} dtype={output.dtype} "
        f"added_parameters={sum(p.numel() for p in model.parameters())}"
    )
    print(
        f"shape={tuple(output.shape)} symmetric={symmetric} "
        f"diagonal_ones={diagonal}"
    )
    print(
        f"reference_match={matches} "
        f"max_output_error={(output - expected).abs().max().item():.3g}"
    )
    for name, (
        finite,
        nonzero,
        inactive_reason,
    ) in gradient_diagnostics.items():
        if inactive_reason:
            print(
                f"{name}_gradient_finite={finite} "
                f"{name}_gradient_nonzero={nonzero} "
                f"{name}_gradient_status=mathematically_inactive "
                f"({inactive_reason})"
            )
        else:
            print(f"{name}_gradient_finite_nonzero={finite and nonzero}")
    if not (
        matches and symmetric and diagonal and all(gradient_checks.values())
    ):
        raise RuntimeError("Numerical reference or backward check failed")


if __name__ == "__main__":
    main()