#!/usr/bin/env python3
"""Correctness tests for transformer-layout-primitives."""

from __future__ import annotations

import argparse
import importlib
import os
import sys
from pathlib import Path

import torch

ROOT = Path(__file__).resolve().parents[2]
PACKAGE = ROOT / "transformer-layout-primitives"
REGISTRATION_INCLUDE = ROOT.parent / "kernels" / "kernel-builder" / "src" / "pyproject" / "templates" / "torch"


class SourceOps:
    def __init__(self, namespace: str) -> None:
        self.ops = getattr(torch.ops, namespace)

    def fill_neginf_bf16(self, dst):
        self.ops.fill_neginf_bf16(dst)
        return dst

    def add_bias_bf16_(self, data, bias):
        self.ops.add_bias_bf16_(data, bias)
        return data

    def repeat_interleave_heads_bf16(self, src, repeat, out=None):
        if out is None:
            out = torch.empty((src.shape[0], src.shape[1] * repeat, src.shape[2]), device=src.device, dtype=src.dtype)
        self.ops.repeat_interleave_heads_bf16(src, int(repeat), out)
        return out

    def text_gather_bf16(self, src, batch, seq, out=None):
        if out is None:
            out = torch.empty((2 * batch, src.shape[1]), device=src.device, dtype=src.dtype)
        self.ops.text_gather_bf16(src, int(batch), int(seq), out)
        return out

    def text_scatter_bf16(self, dst, src, batch, seq):
        self.ops.text_scatter_bf16(dst, src, int(batch), int(seq))
        return dst

    def rope_rotate_half_bf16_(self, x, cos, sin):
        self.ops.rope_rotate_half_bf16_(x, cos, sin)
        return x

    def qk_rmsnorm_rope_bf16_(self, qk, weight, cos, sin, eps=1e-6):
        self.ops.qk_rmsnorm_rope_bf16_(qk, weight, cos, sin, float(eps))
        return qk

    def qk_pair_rmsnorm_rope_bf16(
        self, q, k, q_weight, k_weight, cos, sin, eps=1e-6, q_out=None, k_out=None
    ):
        if q_out is None:
            q_out = torch.empty_like(q)
        if k_out is None:
            k_out = torch.empty_like(k)
        self.ops.qk_pair_rmsnorm_rope_bf16(
            q, k, q_weight, k_weight, cos, sin, float(eps), q_out, k_out
        )
        return q_out, k_out

    def gather_rows_bf16(self, src, row_indices, out=None):
        if out is None:
            out = torch.empty(
                (row_indices.numel(), src.shape[1]), device=src.device, dtype=src.dtype
            )
        self.ops.gather_rows_bf16(src, row_indices, out)
        return out

    def scatter_rows_bf16(self, src, row_indices, destination_rows, out=None):
        if out is None:
            out = torch.zeros(
                (destination_rows, src.shape[1]), device=src.device, dtype=src.dtype
            )
        self.ops.scatter_rows_bf16(src, row_indices, out)
        return out


def _arch_list() -> str:
    major, minor = torch.cuda.get_device_capability(0)
    return "12.0a" if major >= 12 else f"{major}.{minor}"


def load_source_ops() -> SourceOps:
    from torch.utils.cpp_extension import load

    os.environ.setdefault("TORCH_CUDA_ARCH_LIST", _arch_list())
    namespace = "transformer_layout_primitives_source_test"
    load(
        name=namespace,
        sources=[
            str(PACKAGE / "torch-ext" / "torch_binding.cpp"),
            str(PACKAGE / "csrc" / "transformer_layout_primitives.cu"),
        ],
        extra_include_paths=[str(PACKAGE / "csrc"), str(REGISTRATION_INCLUDE)],
        extra_cflags=["-O3", "-DCUDA_KERNEL"],
        extra_cuda_cflags=[
            "-O3",
            "--expt-relaxed-constexpr",
            "-DCUDA_KERNEL",
            "-U__CUDA_NO_BFLOAT16_CONVERSIONS__",
            "-U__CUDA_NO_BFLOAT16_OPERATORS__",
            "-U__CUDA_NO_BFLOAT162_OPERATORS__",
        ],
        is_python_module=False,
        verbose=False,
    )
    return SourceOps(namespace)


def load_installed_ops(artifact: str | None):
    if artifact:
        sys.path.insert(0, artifact)
    try:
        return importlib.import_module("transformer_layout_primitives")
    finally:
        if artifact:
            sys.path.remove(artifact)


def rotate_half_ref(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
    half = x.shape[-1] // 2
    lo = x[..., :half].float()
    hi = x[..., half:].float()
    c = cos[:, None, :half].float()
    s = sin[:, None, :half].float()
    return torch.cat([lo * c - hi * s, hi * c + lo * s], dim=-1).to(torch.bfloat16)


def qk_rmsnorm_rope_ref(
    qk: torch.Tensor,
    weight: torch.Tensor,
    cos: torch.Tensor,
    sin: torch.Tensor,
    eps: float = 1e-6,
) -> torch.Tensor:
    rms = torch.rsqrt((qk.float() * qk.float()).mean(dim=-1, keepdim=True) + eps)
    normed = (qk.float() * rms * weight.float()).to(torch.bfloat16)
    return rotate_half_ref(normed, cos, sin)


def qk_pair_rmsnorm_rope_ref(
    q: torch.Tensor,
    k: torch.Tensor,
    q_weight: torch.Tensor,
    k_weight: torch.Tensor,
    cos: torch.Tensor,
    sin: torch.Tensor,
    eps: float = 1e-6,
) -> tuple[torch.Tensor, torch.Tensor]:
    return (
        qk_rmsnorm_rope_ref(q, q_weight, cos, sin, eps),
        qk_rmsnorm_rope_ref(k, k_weight, cos, sin, eps),
    )


def metrics(got: torch.Tensor, ref: torch.Tensor) -> tuple[float, float, float]:
    diff = (got.float() - ref.float()).abs()
    cos = torch.nn.functional.cosine_similarity(got.float().flatten(), ref.float().flatten(), dim=0).item()
    return float(diff.max().item()), float(diff.mean().item()), float(cos)


def bf16_max_ulp(got: torch.Tensor, ref: torch.Tensor) -> int:
    def ordered_bits(value: torch.Tensor) -> torch.Tensor:
        bits = value.contiguous().view(torch.int16).to(torch.int32) & 0xFFFF
        magnitude = bits & 0x7FFF
        return torch.where(
            (bits & 0x8000) != 0,
            0x8000 - magnitude,
            0x8000 + magnitude,
        )

    return int(
        (ordered_bits(got) - ordered_bits(ref)).abs().max().item()
    )


def assert_close(
    name: str,
    got: torch.Tensor,
    ref: torch.Tensor,
    atol: float,
    cos_min: float,
    rtol: float = 0.0,
) -> None:
    max_abs, mean_abs, cos = metrics(got, ref)
    diff = (got.float() - ref.float()).abs()
    p99_abs = float(torch.quantile(diff, 0.99).item())
    tolerance = atol + rtol * ref.float().abs()
    violations = int((diff > tolerance).sum().item())
    max_ulp = bf16_max_ulp(got, ref)
    print(
        f"{name}: max_abs={max_abs:.6f} p99_abs={p99_abs:.6f} "
        f"mean_abs={mean_abs:.6e} max_ulp={max_ulp} cosine={cos:.8f} "
        f"violations={violations} rtol={rtol} atol={atol}"
    )
    if violations or cos < cos_min:
        raise AssertionError(
            f"{name} failed: max_abs={max_abs} p99_abs={p99_abs} "
            f"max_ulp={max_ulp} cosine={cos} violations={violations}"
        )


def expect_runtime_error(name: str, fn) -> None:
    try:
        fn()
    except RuntimeError:
        print(f"{name}: rejected")
        return
    raise AssertionError(f"{name}: expected RuntimeError")


def run(ops, mode: str) -> int:
    torch.manual_seed(31)
    count = 0
    layout_shapes = [(2, 5, 128), (3, 49, 256)] if mode == "smoke" else [
        (1, 1, 128),
        (2, 5, 128),
        (3, 49, 256),
        (4, 256, 1024),
        (2, 2520, 2048),
    ]
    for batch, seq, dim in layout_shapes:
        x = torch.randn((batch * seq, dim), device="cuda", dtype=torch.bfloat16)

        dst = torch.empty_like(x)
        ops.fill_neginf_bf16(dst)
        ref_neginf = torch.full_like(x, -1e30)
        torch.testing.assert_close(dst.float().cpu(), ref_neginf.float().cpu(), rtol=0, atol=0)
        count += 1

        bias = torch.randn((dim,), device="cuda", dtype=torch.bfloat16)
        got = x.clone()
        ops.add_bias_bf16_(got, bias)
        ref = (x.float() + bias.float()).to(torch.bfloat16)
        torch.testing.assert_close(got.cpu(), ref.cpu(), rtol=0, atol=0)
        count += 1

        gathered = ops.text_gather_bf16(x, batch, seq)
        ref_gather = torch.stack([x[b * seq + offset] for b in range(batch) for offset in (0, seq - 1)], dim=0)
        torch.testing.assert_close(gathered.cpu(), ref_gather.cpu(), rtol=0, atol=0)
        count += 1

        scattered = torch.zeros_like(x)
        ops.text_scatter_bf16(scattered, gathered, batch, seq)
        ref_scatter = torch.zeros_like(x)
        for b in range(batch):
            ref_scatter[b * seq] = gathered[2 * b]
            ref_scatter[b * seq + seq - 1] = gathered[2 * b + 1]
        torch.testing.assert_close(scattered.cpu(), ref_scatter.cpu(), rtol=0, atol=0)
        count += 1

    cpu_rng_state = torch.random.get_rng_state()
    cuda_rng_state = torch.cuda.get_rng_state()
    indexed_shapes = [(17, 5, 64), (277, 51, 128)] if mode == "smoke" else [
        (1, 1, 8),
        (17, 5, 64),
        (277, 51, 128),
        (128, 60, 2048),
        (2520, 105, 1152),
        (5070, 257, 4096),
    ]
    for source_rows, selected_rows, hidden in indexed_shapes:
        src = torch.randn(
            (source_rows, hidden), device="cuda", dtype=torch.bfloat16
        )
        indices = torch.randperm(source_rows, device="cuda", dtype=torch.int64)[
            :selected_rows
        ].contiguous()
        got = ops.gather_rows_bf16(src, indices)
        ref = src.index_select(0, indices)
        torch.testing.assert_close(got, ref, rtol=0, atol=0)
        count += 1

        destination_rows = source_rows + 3
        got = ops.scatter_rows_bf16(ref, indices, destination_rows)
        ref_scatter = torch.zeros(
            (destination_rows, hidden), device="cuda", dtype=torch.bfloat16
        )
        ref_scatter.index_copy_(0, indices, ref)
        torch.testing.assert_close(got, ref_scatter, rtol=0, atol=0)
        count += 1
    torch.random.set_rng_state(cpu_rng_state)
    torch.cuda.set_rng_state(cuda_rng_state)

    repeat_shapes = [(17, 4, 64, 2), (128, 8, 128, 4)] if mode == "smoke" else [
        (1, 1, 64, 8),
        (17, 4, 64, 2),
        (128, 8, 128, 4),
        (2520, 8, 128, 4),
    ]
    for seq, heads, dim, repeat in repeat_shapes:
        src = torch.randn((seq, heads, dim), device="cuda", dtype=torch.bfloat16)
        got = ops.repeat_interleave_heads_bf16(src, repeat)
        ref = src.repeat_interleave(repeat, dim=1)
        torch.testing.assert_close(got.cpu(), ref.cpu(), rtol=0, atol=0)
        count += 1

    rope_shapes = [(17, 4, 64), (128, 8, 128)] if mode == "smoke" else [
        (1, 1, 64),
        (17, 4, 64),
        (128, 8, 128),
        (2520, 32, 128),
    ]
    for seq, heads, dim in rope_shapes:
        x = torch.randn((seq, heads, dim), device="cuda", dtype=torch.bfloat16)
        cos = torch.randn((seq, dim), device="cuda", dtype=torch.bfloat16)
        sin = torch.randn((seq, dim), device="cuda", dtype=torch.bfloat16)
        got = x.clone()
        ops.rope_rotate_half_bf16_(got, cos, sin)
        ref = rotate_half_ref(x, cos, sin)
        assert_close(f"rope seq={seq} heads={heads} dim={dim}", got, ref, atol=0, cos_min=0.999999)
        count += 1

        weight = torch.randn((dim,), device="cuda", dtype=torch.bfloat16)
        got = x.clone()
        ops.qk_rmsnorm_rope_bf16_(got, weight, cos, sin)
        ref = qk_rmsnorm_rope_ref(x, weight, cos, sin)
        # The CUDA kernel is bitwise-gated against the established staged
        # native path below. PyTorch eager may use a different FP32 reduction
        # order, so this semantic comparison uses BF16 elementwise tolerances.
        assert_close(
            f"qk_rmsnorm_rope seq={seq} heads={heads} dim={dim}",
            got,
            ref,
            atol=0.015625,
            rtol=0.02,
            cos_min=0.999999,
        )
        count += 1

    pair_shapes = [(17, 16, 8, 128), (49, 16, 16, 72)] if mode == "smoke" else [
        (1, 16, 8, 128),
        (17, 16, 8, 128),
        (49, 16, 16, 72),
        (51, 16, 16, 80),
        (65, 32, 8, 128),
        (277, 16, 8, 128),
        (512, 24, 24, 128),
        (2520, 24, 24, 128),
        (5070, 24, 24, 128),
    ]
    for rows, q_heads, k_heads, dim in pair_shapes:
        q = torch.randn((rows, q_heads, dim), device="cuda", dtype=torch.bfloat16)
        k = torch.randn((rows, k_heads, dim), device="cuda", dtype=torch.bfloat16)
        q_weight = torch.randn((dim,), device="cuda", dtype=torch.bfloat16)
        k_weight = torch.randn((dim,), device="cuda", dtype=torch.bfloat16)
        angles = torch.randn((rows, dim // 2), device="cuda", dtype=torch.float32)
        cos_half = angles.cos().to(torch.bfloat16)
        sin_half = angles.sin().to(torch.bfloat16)
        cos = torch.cat((cos_half, cos_half), dim=-1)
        sin = torch.cat((sin_half, sin_half), dim=-1)
        got_q, got_k = ops.qk_pair_rmsnorm_rope_bf16(
            q, k, q_weight, k_weight, cos, sin
        )
        staged_q = q.clone()
        staged_k = k.clone()
        ops.qk_rmsnorm_rope_bf16_(staged_q, q_weight, cos, sin)
        ops.qk_rmsnorm_rope_bf16_(staged_k, k_weight, cos, sin)
        torch.testing.assert_close(got_q, staged_q, rtol=0, atol=0)
        torch.testing.assert_close(got_k, staged_k, rtol=0, atol=0)
        ref_q, ref_k = qk_pair_rmsnorm_rope_ref(
            q, k, q_weight, k_weight, cos, sin
        )
        label = f"qk_pair rows={rows} qh={q_heads} kh={k_heads} dim={dim}"
        # The fused path is already required to be bitwise equal to the
        # established staged native kernels above. Torch 2.11's eager
        # reduction order can differ by up to 0.03125 in BF16 at large rows.
        assert_close(
            f"{label}/q", got_q, ref_q, atol=0.015625, rtol=0.02, cos_min=0.999999
        )
        assert_close(
            f"{label}/k", got_k, ref_k, atol=0.015625, rtol=0.02, cos_min=0.999999
        )
        count += 4

    q = torch.randn((17, 4, 64), device="cuda", dtype=torch.bfloat16)
    k = torch.randn((17, 2, 64), device="cuda", dtype=torch.bfloat16)
    weight = torch.ones((64,), device="cuda", dtype=torch.bfloat16)
    cos = torch.ones((17, 64), device="cuda", dtype=torch.bfloat16)
    sin = torch.zeros_like(cos)
    expect_runtime_error(
        "qk_pair mismatched rows",
        lambda: ops.qk_pair_rmsnorm_rope_bf16(
            q, k[:-1].contiguous(), weight, weight, cos, sin
        ),
    )
    expect_runtime_error(
        "qk_pair invalid head_dim",
        lambda: ops.qk_pair_rmsnorm_rope_bf16(
            q[:, :, :-1].contiguous(),
            k[:, :, :-1].contiguous(),
            weight[:-1].contiguous(),
            weight[:-1].contiguous(),
            cos[:, :-1].contiguous(),
            sin[:, :-1].contiguous(),
        ),
    )
    expect_runtime_error(
        "qk_pair noncontiguous",
        lambda: ops.qk_pair_rmsnorm_rope_bf16(
            q.transpose(0, 1), k, weight, weight, cos, sin
        ),
    )
    count += 3

    return count


def run_compile_default_eps(ops) -> int:
    seq, heads, dim = 17, 4, 64
    x = torch.randn((seq, heads, dim), device="cuda", dtype=torch.bfloat16)
    weight = torch.randn((dim,), device="cuda", dtype=torch.bfloat16)
    cos = torch.randn((seq, dim), device="cuda", dtype=torch.bfloat16)
    sin = torch.randn((seq, dim), device="cuda", dtype=torch.bfloat16)
    ref = qk_rmsnorm_rope_ref(x, weight, cos, sin)

    def invoke(qk, rms_weight, rope_cos, rope_sin):
        return ops.qk_rmsnorm_rope_bf16_(qk, rms_weight, rope_cos, rope_sin)

    compiled = torch.compile(invoke, fullgraph=True)
    got = x.clone()
    compiled(got, weight, cos, sin)
    assert_close("qk_rmsnorm_rope compile default eps", got, ref, atol=0.015625, cos_min=0.999999)

    k = torch.randn((seq, 2, dim), device="cuda", dtype=torch.bfloat16)
    k_weight = torch.randn((dim,), device="cuda", dtype=torch.bfloat16)
    pair_ref = qk_pair_rmsnorm_rope_ref(x, k, weight, k_weight, cos, sin)

    def invoke_pair(q, key, qw, kw, rope_cos, rope_sin):
        return ops.qk_pair_rmsnorm_rope_bf16(q, key, qw, kw, rope_cos, rope_sin)

    compiled_pair = torch.compile(invoke_pair, fullgraph=True)
    got_q, got_k = compiled_pair(x, k, weight, k_weight, cos, sin)
    assert_close("qk_pair compile default eps/q", got_q, pair_ref[0], atol=0.015625, cos_min=0.999999)
    assert_close("qk_pair compile default eps/k", got_k, pair_ref[1], atol=0.015625, cos_min=0.999999)
    return 3


def run_indexed_compile(ops) -> int:
    src = torch.randn((277, 128), device="cuda", dtype=torch.bfloat16)
    indices = torch.randperm(277, device="cuda", dtype=torch.int64)[:51].contiguous()

    def invoke(src, indices):
        gathered = ops.gather_rows_bf16(src, indices)
        return ops.scatter_rows_bf16(gathered, indices, src.shape[0])

    eager = invoke(src, indices)
    compiled = torch.compile(invoke, fullgraph=True)(src, indices)
    torch.testing.assert_close(compiled, eager, rtol=0, atol=0)
    print("PASS gather/scatter torch.compile fullgraph")
    return 1


def run_cosmos_edge_indexed_graph(ops) -> int:
    source_rows, selected_rows, hidden = 128, 60, 2048
    src = torch.randn((source_rows, hidden), device="cuda", dtype=torch.bfloat16)
    indices = torch.randperm(
        source_rows, device="cuda", dtype=torch.int64
    )[:selected_rows].contiguous()
    gathered = torch.empty(
        (selected_rows, hidden), device="cuda", dtype=torch.bfloat16
    )
    scattered = torch.zeros(
        (source_rows, hidden), device="cuda", dtype=torch.bfloat16
    )

    graph = torch.cuda.CUDAGraph()
    torch.cuda.synchronize()
    with torch.cuda.graph(graph):
        ops.gather_rows_bf16(src, indices, out=gathered)
        ops.scatter_rows_bf16(gathered, indices, source_rows, out=scattered)
    graph.replay()
    torch.cuda.synchronize()
    expected_gathered = src.index_select(0, indices)
    expected_scattered = torch.zeros_like(scattered)
    expected_scattered.index_copy_(0, indices, expected_gathered)
    torch.testing.assert_close(gathered, expected_gathered, rtol=0.0, atol=0.0)
    torch.testing.assert_close(scattered, expected_scattered, rtol=0.0, atol=0.0)
    first = (gathered.clone(), scattered.clone())
    graph.replay()
    torch.cuda.synchronize()
    torch.testing.assert_close(gathered, first[0], rtol=0.0, atol=0.0)
    torch.testing.assert_close(scattered, first[1], rtol=0.0, atol=0.0)
    print("PASS Cosmos3-Edge gather/scatter CUDA Graph replay")
    return 2


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--backend", choices=["source", "installed"], default="source")
    parser.add_argument("--artifact", default=None)
    parser.add_argument("--mode", choices=["smoke", "full"], default="smoke")
    args = parser.parse_args()
    ops = load_source_ops() if args.backend == "source" else load_installed_ops(args.artifact)
    count = run(ops, args.mode)
    count += run_indexed_compile(ops)
    if args.mode == "full":
        count += run_cosmos_edge_indexed_graph(ops)
    if args.backend == "installed":
        count += run_compile_default_eps(ops)
    print(f"transformer-layout-primitives {args.backend} {args.mode}: passed {count}/{count}")
    return 0


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