"""Parity + integration tests for the OrbitQuant Triton kernels.

Strategy: compare the Triton kernels against the pure-PyTorch reference in
_orbitquant.py (which is the numerical ground truth and the converter-matched
implementation).  All tests require CUDA + Triton; skipped otherwise.
"""

from __future__ import annotations

import pytest
import torch

from _intcrush import load

_oq = load("_orbitquant")
_qu = load("_quant_utils")
_tr = load("_triton_runtime")

CUDA = torch.cuda.is_available() and _tr.TRITON_AVAILABLE
TRITON_RPBH = CUDA and _tr.TRITON_RPBH_ROTATE
TRITON_CB = CUDA and _tr.TRITON_CODEBOOK_DEQUANT

# Realistic (d, h) pairs from the paper / FLUX / Wan / CogVideoX.
_RPBH_SHAPES = [(3072, 1024), (12288, 4096), (256, 256), (1920, 64), (4096, 4096)]
# A (d, h) where d_padded > d exercises the zero-padding path.
_PAD_SHAPES = [(3000, 1024)]   # d_padded = ceil(3000/1024)*1024 = 3072


def _make_codebook(bits: int) -> torch.Tensor:
    """Uniform codebook (tests only -- production uses Lloyd-Max on f_d)."""
    n = 1 << bits
    return torch.linspace(-1.0, 1.0, n, dtype=torch.float32)


# ── RPBH ────────────────────────────────────────────────────────────────────

@pytest.mark.skipif(not TRITON_RPBH, reason="CUDA/Triton RPBH unavailable")
@pytest.mark.parametrize("d,h", _RPBH_SHAPES + _PAD_SHAPES)
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
def test_rpbh_parity_vs_pytorch(d, h, dtype):
    perm, signs, _ = _oq.make_rpbh_params(d, seed=42)
    d_padded = perm.shape[0]
    perm = perm.cuda()
    signs = signs.cuda()
    M = 64
    x = torch.randn(M, d, dtype=dtype, device="cuda")

    ref = _oq.rpbh_rotate(x, perm, signs, h)
    got = _tr.rpbh_rotate_triton(x, perm, signs, h, d_padded)

    assert got.shape == (M, d_padded), f"shape {got.shape} != ({M},{d_padded})"
    assert got.dtype == x.dtype
    # fp32 internal compute; tolerate fp16/bf16 storage rounding.
    atol = 1e-3 if dtype != torch.float32 else 1e-4
    rtol = 1e-3
    torch.testing.assert_close(got.float(), ref.float(), atol=atol, rtol=rtol)


@pytest.mark.skipif(not TRITON_RPBH, reason="CUDA/Triton RPBH unavailable")
@pytest.mark.parametrize("d,h", _RPBH_SHAPES)
def test_rpbh_norm_preservation(d, h):
    perm, signs, _ = _oq.make_rpbh_params(d, seed=7)
    perm = perm.cuda()
    signs = signs.cuda()
    x = torch.randn(32, d, device="cuda")
    y = _tr.rpbh_rotate_triton(x, perm, signs, h)
    # RPBH is orthogonal: ||y|| == ||x|| (up to padding zeros, which x lacks).
    rel = (x.norm() - y.norm()).abs().item() / x.norm().item()
    assert rel < 1e-3, f"norm not preserved: rel={rel}"


@pytest.mark.skipif(not TRITON_RPBH, reason="CUDA/Triton RPBH unavailable")
@pytest.mark.parametrize("d,h", _RPBH_SHAPES)
def test_rpbh_cancellation(d, h):
    """(W @ Pi^T) @ (x @ Pi^T)^T == W @ x^T  (the sacred contract)."""
    perm, signs, _ = _oq.make_rpbh_params(d, seed=42)
    perm = perm.cuda()
    signs = signs.cuda()
    W = torch.randn(16, d, device="cuda", dtype=torch.float32)
    x = torch.randn(8, d, device="cuda", dtype=torch.float32)

    W_rot = _tr.rpbh_rotate_triton(W, perm, signs, h)
    x_rot = _tr.rpbh_rotate_triton(x, perm, signs, h)

    y_rot = W_rot @ x_rot.t()
    y_ref = W @ x.t()
    rel = (y_rot - y_ref).norm().item() / y_ref.norm().item()
    assert rel < 1e-3, f"cancellation broken: rel={rel}"


@pytest.mark.skipif(not TRITON_RPBH, reason="CUDA/Triton RPBH unavailable")
def test_rpbh_dtype_preservation():
    d, h = 1024, 1024
    perm, signs, _ = _oq.make_rpbh_params(d, seed=1)
    perm = perm.cuda()
    signs = signs.cuda()
    for dt in (torch.float16, torch.bfloat16, torch.float32):
        x = torch.randn(4, d, dtype=dt, device="cuda")
        y = _tr.rpbh_rotate_triton(x, perm, signs, h)
        assert y.dtype == dt


@pytest.mark.skipif(not TRITON_RPBH, reason="CUDA/Triton RPBH unavailable")
def test_rpbh_zero_padding():
    """d < d_padded: padded positions contribute 0, matching rpbh_rotate."""
    d, h = 3000, 1024          # d_padded = 3072
    perm, signs, _ = _oq.make_rpbh_params(d, seed=42)
    d_padded = perm.shape[0]
    assert d_padded > d
    perm = perm.cuda()
    signs = signs.cuda()
    x = torch.randn(16, d, device="cuda", dtype=torch.float32)
    ref = _oq.rpbh_rotate(x, perm, signs, h)
    got = _tr.rpbh_rotate_triton(x, perm, signs, h, d_padded)
    torch.testing.assert_close(got, ref, atol=1e-4, rtol=1e-4)


# ── Codebook dequant ──────────────────────────────────────────────────────────

@pytest.mark.skipif(not TRITON_CB, reason="CUDA/Triton codebook dequant unavailable")
@pytest.mark.parametrize("bits", [4, 8])
def test_codebook_dequant_parity(bits):
    d = 256
    N = 32
    codebook = _make_codebook(bits).cuda()
    W = torch.randn(N, d, device="cuda") * 0.1
    norms = W.norm(dim=1, keepdim=True).clamp(min=1e-10)
    W_dir = W / norms

    # Nearest-centroid indices.
    boundaries = (codebook[:-1] + codebook[1:]) / 2.0
    indices = torch.searchsorted(boundaries, W_dir.reshape(-1)).reshape(N, d)
    indices = indices.clamp(0, (1 << bits) - 1)
    offset = 1 << (bits - 1)
    indices_i8 = (indices - offset).to(torch.int8)

    if bits == 4:
        packed = _qu.pack_int4(indices_i8.cpu()).to(torch.uint8).cuda()
    else:
        packed = indices_i8.cuda()

    norms_bf = norms.to(torch.bfloat16).cuda()

    # Reference returns float32.
    if bits == 4:
        ref = _oq.codebook_dequant_int4(packed, norms_bf, codebook, d, d)
    else:
        ref = _oq.codebook_dequant_int8(packed, norms_bf, codebook, d, d)

    # Triton default -> fp16; compare in fp32 with fp16 tolerance.
    got = _tr.codebook_dequant_triton(packed, norms_bf, codebook, d, bits)
    assert got.shape == (N, d)
    torch.testing.assert_close(got.float(), ref.float(), atol=2e-3, rtol=1e-2)

    # fp32 output path (used by OrbitQuantLayout.dequantize for LoRA parity).
    out32 = torch.empty(N, d, dtype=torch.float32, device="cuda")
    got32 = _tr.codebook_dequant_triton(packed, norms_bf, codebook, d, bits, out=out32)
    torch.testing.assert_close(got32, ref, atol=1e-5, rtol=1e-5)


@pytest.mark.skipif(not TRITON_CB, reason="CUDA/Triton codebook dequant unavailable")
def test_codebook_dequant_no_trim():
    """Output last dim is d_padded, never trimmed to d (inference-basis match)."""
    bits = 4
    d_padded = 256
    N = 8
    codebook = _make_codebook(bits).cuda()
    packed = torch.zeros(N, (d_padded + 1) // 2, dtype=torch.uint8, device="cuda")
    norms = torch.ones(N, dtype=torch.bfloat16, device="cuda")
    out = _tr.codebook_dequant_triton(packed, norms, codebook, d_padded, bits)
    assert out.shape == (N, d_padded), f"shape {out.shape} != ({N},{d_padded})"
    assert out.dtype == torch.float16


# ── Dispatch / fallback ─────────────────────────────────────────────────────

@pytest.mark.skipif(not CUDA, reason="CUDA unavailable")
def test_rpbh_fallback_when_triton_disabled(monkeypatch):
    """With TRITON_RPBH_ROTATE=False, the fast dispatcher uses PyTorch."""
    if not _tr.TRITON_RPBH_ROTATE:
        pytest.skip("needs Triton RPBH loaded to test the flag-off path")
    d, h = 256, 256
    perm, signs, _ = _oq.make_rpbh_params(d, seed=42)
    perm = perm.cuda()
    signs = signs.cuda()
    x = torch.randn(8, d, device="cuda", dtype=torch.float32)

    monkeypatch.setattr(_tr, "TRITON_RPBH_ROTATE", False)
    y = _oq.rpbh_rotate_fast(x, perm, signs, h)
    ref = _oq.rpbh_rotate(x, perm, signs, h)
    torch.testing.assert_close(y, ref, atol=1e-4, rtol=1e-4)


@pytest.mark.skipif(not CUDA, reason="CUDA unavailable")
def test_codebook_dequant_fallback_when_triton_disabled(monkeypatch):
    """With TRITON_CODEBOOK_DEQUANT=False, the fast dispatcher uses PyTorch."""
    if not _tr.TRITON_CODEBOOK_DEQUANT:
        pytest.skip("needs Triton codebook dequant loaded to test the flag-off path")
    bits = 4
    d = 256
    codebook = _make_codebook(bits).cuda()
    indices_i8 = torch.zeros(16, d, dtype=torch.int8, device="cuda")
    packed = _qu.pack_int4(indices_i8.cpu()).to(torch.uint8).cuda()
    norms = torch.ones(16, dtype=torch.bfloat16, device="cuda")

    monkeypatch.setattr(_tr, "TRITON_CODEBOOK_DEQUANT", False)
    y = _oq.codebook_dequant_fast(packed, norms, codebook, d, d, bits)
    ref = _oq.codebook_dequant(packed, norms, codebook, d, d, bits)
    torch.testing.assert_close(y.float(), ref.float(), atol=1e-5, rtol=1e-5)


@pytest.mark.skipif(not CUDA, reason="CUDA unavailable")
def test_rpbh_cpu_uses_pytorch_reference():
    """CPU inputs always go through the PyTorch reference (no Triton on CPU)."""
    d, h = 256, 256
    perm, signs, _ = _oq.make_rpbh_params(d, seed=42)
    x = torch.randn(4, d, dtype=torch.float32)  # CPU
    y = _oq.rpbh_rotate_fast(x, perm, signs, h)
    ref = _oq.rpbh_rotate(x, perm, signs, h)
    torch.testing.assert_close(y, ref, atol=1e-5, rtol=1e-5)


# ── Integration: forward-path contracts ───────────────────────────────────────

@pytest.mark.skipif(not TRITON_RPBH, reason="CUDA/Triton RPBH unavailable")
def test_adaln_never_rotated():
    """AdaLN-RTN activations must NOT be RPBH-rotated (paper Sec B.2).

    An adaln_rtn layer has _intcrush_oq_perm = None (loader sets RPBH only
    for codebook layers; adaln gets a {"perm": None, ...} fallback).
    _rpbh_rotate would raise on None perm, which is exactly the guard --
    the forward path calls adaln_rtn_dequant + F.linear directly, bypassing
    _rpbh_rotate.  This test documents that contract.
    """
    d, h = 256, 1024
    perm, signs, _ = _oq.make_rpbh_params(d, seed=42)
    perm = perm.cuda()
    signs = signs.cuda()
    x = torch.randn(8, d, device="cuda", dtype=torch.float32)

    # A codebook-layer activation, rotated, is well-defined:
    y = _tr.rpbh_rotate_triton(x, perm, signs, h)
    assert y.shape == (8, perm.shape[0])

    # An adaln-layer activation is used RAW -- this is the reference value:
    y_adaln = x   # literally x_2d unchanged
    assert y_adaln.shape == (8, d)


@pytest.mark.skipif(not (TRITON_RPBH and TRITON_CB), reason="needs both kernels")
def test_forward_orbitquant_matches_reference():
    """End-to-end: rpbh(x) @ dequant(W) == reference (cancellation preserves Wx)."""
    d, h = 1024, 1024
    N_out = 64
    M = 16
    bits = 4
    codebook = _make_codebook(bits).cuda()

    # Build a random weight, rotate it offline (converter-side), quantize to codebook.
    W = torch.randn(N_out, d, device="cuda", dtype=torch.float32)

    perm, signs, _ = _oq.make_rpbh_params(d, seed=42)
    perm = perm.cuda()
    signs = signs.cuda()
    d_padded = perm.shape[0]

    # Rotate weight into the shared basis (offline step the converter does).
    W_rot = _oq.rpbh_rotate(W, perm, signs, h)

    # Quantize the rotated directions to the codebook (simulates converter).
    norms = W_rot.norm(dim=1, keepdim=True).clamp(min=1e-10)
    Wrot_dir = W_rot / norms
    boundaries = (codebook[:-1] + codebook[1:]) / 2.0
    indices = torch.searchsorted(boundaries, Wrot_dir.reshape(-1)).reshape(N_out, d_padded)
    indices = indices.clamp(0, (1 << bits) - 1)
    offset = 1 << (bits - 1)
    indices_i8 = (indices - offset).to(torch.int8)
    packed = _qu.pack_int4(indices_i8.cpu()).to(torch.uint8).cuda()
    norms_bf = norms.to(torch.bfloat16).cuda()

    # Reference forward (pure PyTorch):
    x = torch.randn(M, d, device="cuda", dtype=torch.float32)
    x_rot_ref = _oq.rpbh_rotate(x, perm, signs, h)
    W_f_ref = _oq.codebook_dequant_int4(packed, norms_bf, codebook, d, d_padded)
    y_ref = x_rot_ref @ W_f_ref.to(torch.float32).t()

    # Triton forward:
    x_rot = _tr.rpbh_rotate_triton(x, perm, signs, h, d_padded)
    W_f = _tr.codebook_dequant_triton(
        packed, norms_bf, codebook, d_padded, bits,
        out=torch.empty(N_out, d_padded, dtype=torch.float32, device="cuda"),
    )
    y = x_rot.float() @ W_f.t()

    rel = (y - y_ref).norm().item() / max(y_ref.norm().item(), 1e-12)
    assert rel < 2e-3, f"end-to-end mismatch: rel={rel}"