# examples/hilbert_attention.py
"""Seeded image-token attention with Hilbert ordering and native FlexAttention.

This is an attention-core experiment, not a pretrained HWT classifier or an
ImageNet benchmark reproduction. HSA/HNA use new one-dimensional neighborhoods;
they generally differ from the conventional two-dimensional spatial masks.
Importing this module does not run benchmarks, compile kernels, or download data.

Traversal adapted from the authors' hwt.py:
https://github.com/Yunge6666/Hilbert-Local-Attention/blob/main/hwt.py
FlexAttention API: https://docs.pytorch.org/docs/2.12/nn.attention.flex_attention.html

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
from functools import lru_cache
import math
import statistics
import time

import torch


def _positive_integer(value, name):
    if isinstance(value, bool) or not isinstance(value, int) or value <= 0:
        raise ValueError(f"{name} must be a positive integer")


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: order
@lru_cache(maxsize=128, typed=True)
def hilbert_order(height, width):
    """Return cached sequence index -> row-major flat index for any rectangle."""
    _positive_integer(height, "height")
    _positive_integer(width, "width")
    path = []
    if width >= height:
        _generate2d(0, 0, width, 0, 0, height, path)
    else:
        _generate2d(0, 0, 0, height, width, 0, path)
    return tuple(row * width + col for col, row in path)


# endregion: order


def _validate_token_tensor(tensor, n):
    if not isinstance(tensor, torch.Tensor) or tensor.ndim != 4:
        raise ValueError(
            "Expected a tensor [batch, heads, height*width, head_dim]"
        )
    if tensor.shape[-2] != n or any(
        dimension <= 0 for dimension in tensor.shape
    ):
        raise ValueError(
            "Tensor dimensions must be positive; token count must match height*width"
        )


# region: reorder
def reorder_tokens(tensor, height, width):
    """Move row-major Q/K/V tokens into Hilbert sequence order."""
    order = hilbert_order(height, width)
    _validate_token_tensor(tensor, len(order))
    indices = torch.tensor(order, device=tensor.device, dtype=torch.long)
    return tensor.index_select(-2, indices)


# endregion: reorder


# region: restore
def restore_order(tensor, height, width):
    """Invert the permutation on the token axis, including its gradients."""
    order = hilbert_order(height, width)
    _validate_token_tensor(tensor, len(order))
    indices = torch.tensor(order, device=tensor.device, dtype=torch.long)
    inverse = torch.argsort(indices)
    return tensor.index_select(-2, inverse)


# endregion: restore


# region: window
def window_predicate(n, window_side, shift=0):
    """Build HWA in original Hilbert indices, guarding rotated wraparound."""
    area = window_side**2

    def allowed(i, j):
        a, b = (i + shift) % n, (j + shift) % n
        same_window = (a // area) == (b // area)
        same_wrap = ((i + shift) // n) == ((j + shift) // n)
        return same_window & same_wrap

    return allowed


# endregion: window


# region: neighborhood
def neighborhood_predicate(n, kernel_side):
    """Build HNA with a fixed area neighborhood shifted inward at the ends."""
    area = kernel_side**2
    half = (area - 1) // 2

    def allowed(i, j):
        start = i - half
        if isinstance(start, torch.Tensor):
            start = start.clamp(0, n - area)
        else:
            start = min(max(start, 0), n - area)
        return (j >= start) & (j < start + area)

    return allowed


# endregion: neighborhood


def make_predicate(n, pattern, size, shift=0):
    """Validate geometry and return an allowed-pair scalar/tensor predicate.

    Args:
        n: Number of sequence tokens.
        pattern: "window" (HWA), "slide" (HSA), or "neighborhood" (HNA).
        size: Window side, or an odd positive kernel side for HSA/HNA.
        shift: Integer in [0, n); available only for HWA.

    Returns:
        Callable accepting query/key indices. Boolean True means allowed.
    """
    _positive_integer(n, "n")
    _positive_integer(size, "size")
    if (
        isinstance(shift, bool)
        or not isinstance(shift, int)
        or not 0 <= shift < n
    ):
        raise ValueError("shift must be an integer in [0, n)")
    if pattern not in ("window", "slide", "neighborhood"):
        raise ValueError("pattern must be window, slide, or neighborhood")
    if pattern != "window" and shift:
        raise ValueError("shift is only supported for window attention")
    if pattern != "window" and size % 2 == 0:
        raise ValueError("HSA/HNA kernel side must be odd")
    if pattern == "neighborhood" and size**2 > n:
        raise ValueError("HNA kernel area must not exceed token count")
    if pattern == "window":
        return window_predicate(n, size, shift)
    if pattern == "neighborhood":
        return neighborhood_predicate(n, size)
    half = (size**2 - 1) // 2
    return lambda i, j: abs(i - j) <= half


def row_major_window_predicate(height, width, size):
    """Conventional 2D WSA comparator; edge windows may be partial."""
    _positive_integer(height, "height")
    _positive_integer(width, "width")
    _positive_integer(size, "size")
    return lambda i, j: (
        ((i // width) // size == (j // width) // size)
        & ((i % width) // size == (j % width) // size)
    )


def dense_mask(predicate, n, device="cpu"):
    """Evaluate the same predicate used by FlexAttention as a boolean mask."""
    _positive_integer(n, "n")
    indices = torch.arange(n, device=device)
    return predicate(indices[:, None], indices[None, :])


def _block_kind(valid):
    hits = int(valid.sum())
    if hits == 0:
        return "empty"
    return "full" if hits == valid.numel() else "partial"


# region: blocks
def block_statistics(predicate, n, block_size):
    """Classify valid cells only; empty_fraction counts blocks, not elements."""
    _positive_integer(n, "n")
    _positive_integer(block_size, "block_size")
    counts = {"empty": 0, "partial": 0, "full": 0}
    for q in range(0, n, block_size):
        rows = torch.arange(q, min(q + block_size, n))[:, None]
        for k in range(0, n, block_size):
            cols = torch.arange(k, min(k + block_size, n))[None, :]
            valid = predicate(rows, cols)
            counts[_block_kind(valid)] += 1
    total = math.ceil(n / block_size) ** 2
    return {
        "counts": counts,
        "total": total,
        "empty_fraction": counts["empty"] / total,
    }


# endregion: blocks


def _validate_qkv(q, k, v):
    if not isinstance(q, torch.Tensor) or q.ndim != 4:
        raise ValueError(
            "Q/K/V must have shape [batch, heads, tokens, head_dim]"
        )
    _validate_token_tensor(q, q.shape[-2])
    if not q.is_floating_point():
        raise ValueError("Q/K/V must have floating point dtype")
    for tensor in (k, v):
        if not isinstance(tensor, torch.Tensor) or tensor.shape != q.shape:
            raise ValueError("Q/K/V must have identical shapes")
        if tensor.device != q.device or tensor.dtype != q.dtype:
            raise ValueError("Q/K/V must have identical devices and dtypes")


def dense_attention(q, k, v, mask):
    """Independent reference: softmax(masked QK^T/sqrt(head_dim)) V."""
    _validate_qkv(q, k, v)
    n = q.shape[-2]
    if (
        mask.shape != (n, n)
        or mask.dtype != torch.bool
        or mask.device != q.device
    ):
        raise ValueError(
            "Mask must be bool [tokens, tokens] on the Q/K/V device"
        )
    if not bool(mask.any(-1).all()):
        raise ValueError("Every query must have at least one allowed key")
    scores = (q @ k.transpose(-2, -1)) / math.sqrt(q.shape[-1])
    probabilities = scores.masked_fill(~mask, float("-inf")).softmax(-1)
    return probabilities @ v


# region: flex
def build_flex_attention(predicate, n, device):
    """Construct one native block mask and one reusable compiled function."""
    from torch.nn.attention.flex_attention import create_block_mask
    from torch.nn.attention.flex_attention import flex_attention

    def mask_mod(batch, head, i, j):
        return (i < n) & (j < n) & predicate(i, j)

    block_mask = create_block_mask(
        mask_mod,
        B=None,
        H=None,
        Q_LEN=n,
        KV_LEN=n,
        device=str(device),
        BLOCK_SIZE=128,
    )
    compiled = torch.compile(flex_attention)

    def attend(q, k, v):
        return compiled(q, k, v, block_mask=block_mask)

    return attend


# endregion: flex


class HilbertAttention:
    """Reusable row-major input/output attention with geometry prepared once.

    Args:
        height: Feature map height in tokens.
        width: Feature map width in tokens.
        pattern: window, slide, or neighborhood.
        size: Window or kernel side.
        shift: Positive 1D HWA rotation with wraparound connections masked.
        backend: dense (CPU/CUDA) or flex (CUDA only).
        device: Tensor device. CUDA must be available when selected.
        block_size: Diagnostic block size; Flex requires 128.
    """

    def __init__(
        self,
        height,
        width,
        pattern="window",
        size=4,
        shift=0,
        backend="dense",
        device="cpu",
        block_size=128,
    ):
        self.height, self.width = height, width
        order = hilbert_order(height, width)
        self.n = len(order)
        self.predicate = make_predicate(self.n, pattern, size, shift)
        _positive_integer(block_size, "block_size")
        if backend not in ("dense", "flex"):
            raise ValueError("backend must be dense or flex")
        self.device = torch.device(device)
        if self.device.type not in ("cpu", "cuda"):
            raise ValueError("device must be cpu or cuda")
        if backend == "flex" and (
            self.device.type != "cuda" or block_size != 128
        ):
            raise ValueError("FlexAttention requires CUDA and 128-token blocks")
        if self.device.type == "cuda":
            if not torch.cuda.is_available():
                raise RuntimeError(
                    "CUDA requested, but this PyTorch installation has no available CUDA device"
                )
            if self.device.index is None:
                self.device = torch.device("cuda", torch.cuda.current_device())
        self.backend = backend
        self.mask = None
        if backend == "dense":
            self.mask = dense_mask(self.predicate, self.n, self.device)
            self.attend = lambda q, k, v: dense_attention(q, k, v, self.mask)
        else:
            self.attend = build_flex_attention(
                self.predicate, self.n, self.device
            )

    def __call__(self, q, k, v):
        """Reorder Q/K/V, execute the chosen backend, and restore output."""
        _validate_qkv(q, k, v)
        _validate_token_tensor(q, self.n)
        if q.device != self.device:
            raise ValueError(
                "Q/K/V device must match the prepared attention device"
            )
        if self.backend == "flex" and (
            q.dtype not in (torch.float16, torch.bfloat16, torch.float32)
            or q.shape[-1] < 16
            or q.shape[-1] > 128
            or q.shape[-1] & (q.shape[-1] - 1)
        ):
            raise ValueError(
                "Flex requires float16/bfloat16/float32 and power-of-two head_dim in [16, 128]"
            )
        # Public helpers are the actual executed reorder/restore excerpt bodies.
        qh, kh, vh = (
            reorder_tokens(t, self.height, self.width) for t in (q, k, v)
        )
        try:
            output = self.attend(qh, kh, vh)
        except RuntimeError as error:
            if self.backend == "flex":
                raise RuntimeError(
                    "Native CUDA FlexAttention failed; no dense fallback was run"
                ) from error
            raise
        return restore_order(output, self.height, self.width)


def benchmark(attend, q, k, v, warmup=5, iterations=20):
    """Median forward latency including reorder/restore, excluding setup."""
    _positive_integer(warmup, "warmup")
    _positive_integer(iterations, "iterations")

    def synchronize():
        if q.device.type == "cuda":
            torch.cuda.synchronize(q.device)

    with torch.no_grad():
        for _ in range(warmup):
            attend(q, k, v)
        synchronize()
        samples = []
        for _ in range(iterations):
            synchronize()
            start = time.perf_counter()
            attend(q, k, v)
            synchronize()
            samples.append((time.perf_counter() - start) * 1000)
    return statistics.median(samples)


def main(argv=None):
    """Run a bounded correctness experiment, optionally timing forward runs."""
    parser = argparse.ArgumentParser(
        description=__doc__.split("\n\n", maxsplit=1)[0]
    )
    parser.add_argument("--device", choices=("cpu", "cuda"), default="cpu")
    parser.add_argument("--backend", choices=("dense", "flex"), default="dense")
    parser.add_argument(
        "--pattern",
        choices=("window", "slide", "neighborhood"),
        default="window",
    )
    parser.add_argument("--height", type=int, default=8)
    parser.add_argument("--width", type=int, default=8)
    parser.add_argument("--size", type=int, default=4)
    parser.add_argument("--shift", type=int, default=0)
    parser.add_argument("--block-size", type=int, default=128)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--benchmark", action="store_true")
    parser.add_argument("--warmup", type=int, default=5)
    parser.add_argument("--iterations", type=int, default=20)
    args = parser.parse_args(argv)
    if args.height <= 0 or args.width <= 0 or args.height * args.width > 4096:
        parser.error(
            "height and width must be positive with at most 4096 tokens for bounded dense diagnostics"
        )
    if (
        args.block_size <= 0
        or math.ceil(args.height * args.width / args.block_size) > 128
    ):
        parser.error(
            "block-size must be positive and produce at most 128 blocks per axis"
        )
    if (
        args.warmup <= 0
        or args.iterations <= 0
        or args.warmup > 1000
        or args.iterations > 1000
    ):
        parser.error("warmup and iterations must be in [1, 1000]")
    try:
        runner = HilbertAttention(
            args.height,
            args.width,
            args.pattern,
            args.size,
            args.shift,
            args.backend,
            args.device,
            args.block_size,
        )
    except (ValueError, RuntimeError) as error:
        parser.error(str(error))
    generator = torch.Generator().manual_seed(args.seed)
    q, k, v = [
        torch.randn(1, 2, runner.n, 32, generator=generator).to(runner.device)
        for _ in range(3)
    ]
    print(
        f"Configuration: device={args.device}, backend={args.backend}, pattern={args.pattern}, "
        f"height={args.height}, width={args.width}, size={args.size}, shift={args.shift}, "
        f"block_size={args.block_size}, seed={args.seed}, batch=1, heads=2, head_dim=32"
    )
    print(
        "Scope: seeded image-token attention core; no pretrained model, dataset, or ImageNet accuracy."
    )
    with torch.no_grad():
        output = runner(q, k, v)
        print(f"Output shape: {tuple(output.shape)}")
        if args.backend == "flex":
            reference_runner = HilbertAttention(
                args.height,
                args.width,
                args.pattern,
                args.size,
                args.shift,
                "dense",
                args.device,
                args.block_size,
            )
            reference = reference_runner(q, k, v)
            error = float((output - reference).abs().max())
            print(
                f"Dense/Flex maximum absolute error (same Hilbert mask): {error:.8g}"
            )
            torch.testing.assert_close(output, reference, rtol=2e-4, atol=2e-4)
    stats = block_statistics(runner.predicate, runner.n, args.block_size)
    print(
        f"Hilbert blocks: {stats['counts']}; empty blocks={stats['empty_fraction']:.3%} ({stats['total']} total)"
    )
    if args.pattern == "window":
        comparator = row_major_window_predicate(
            args.height, args.width, args.size
        )
        row_stats = block_statistics(comparator, runner.n, args.block_size)
        print(
            f"Unshifted row-major 2D window blocks: {row_stats['counts']}; "
            f"empty blocks={row_stats['empty_fraction']:.3%}"
        )
    else:
        print(
            "HSA/HNA are 1D Hilbert neighborhoods, generally different from 2D row-major neighborhoods."
        )
    if args.benchmark:
        median = benchmark(runner, q, k, v, args.warmup, args.iterations)
        print(
            f"Median forward latency: {median:.4f} ms ({args.iterations} iterations, {args.warmup} warmups)."
        )
        print(
            "Timing includes reorder, attention and restore; excludes compilation/block-mask construction. "
            "GPU runs synchronized individually. This is not a paper speedup claim."
        )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())