import pytest
import torch
import torch.nn.functional as F

from flashinfer import autotune, mm_bf16
from flashinfer.gemm.gemm_base import CUDNN_AVAILABLE
from flashinfer.gemm import is_cuda_tile_available
from flashinfer.utils import get_compute_capability, is_sm100a_supported


@pytest.mark.parametrize("m", [1, 8, 16, 32, 64])
@pytest.mark.parametrize("n", [1024, 2048, 4096])
@pytest.mark.parametrize("k", [1024, 2048, 3072])
@pytest.mark.parametrize("res_dtype", [torch.bfloat16, torch.float16, torch.float32])
@pytest.mark.parametrize("enable_bias", [True, False])
@pytest.mark.parametrize("pdl", [True, False])
@pytest.mark.parametrize(
    "backend",
    ["cudnn", "cutlass", "tgv", "cublaslt", "tinygemm", "cutile", "cute-dsl", "auto"],
)
@pytest.mark.parametrize("auto_tuning", [False, True])
def test_mm_bf16(
    m: int,
    n: int,
    k: int,
    res_dtype: torch.dtype,
    enable_bias: bool,
    pdl: bool,
    backend: str,
    auto_tuning: bool,
):
    compute_capability = get_compute_capability(torch.device(device="cuda"))
    compute_capability_number = compute_capability[0] * 10 + compute_capability[1]
    if not mm_bf16.is_compute_capability_supported(compute_capability_number):
        pytest.skip(
            f"mm_bf16 not supported on current compute capability."
            f"Detected sm{compute_capability_number}."
        )
    if backend != "auto":
        if not mm_bf16.is_backend_supported(backend, compute_capability_number):
            pytest.skip(
                f"{backend} backend not supported on current compute capability."
            )

    if backend == "cudnn" and not CUDNN_AVAILABLE:
        pytest.skip("cuDNN is not available on this system.")

    if backend == "cutile":
        if not is_cuda_tile_available():
            pytest.skip(
                "cuda-tile / tileiras compiler not available in this environment."
            )

    if backend == "cute-dsl":
        if not is_sm100a_supported(torch.device("cuda")):
            pytest.skip("CuTeDSL low-M backend requires SM100/SM103 with CUDA 12.8+.")

        from flashinfer.cute_dsl.utils import is_cute_dsl_available

        if not is_cute_dsl_available():
            pytest.skip("nvidia-cutlass-dsl is not available.")
        if m > 32:
            pytest.skip("CuTeDSL low-M backend requires M <= 32.")
        if res_dtype != torch.bfloat16:
            pytest.skip("CuTeDSL low-M backend requires BF16 output.")

    if backend == "auto" and (enable_bias or pdl):
        pytest.skip("mm_bf16 with auto backend does not support bias or pdl arguments.")

    if backend == "cutlass" and (enable_bias or pdl):
        pytest.skip(
            "mm_bf16 with CUTLASS backend does not support bias or pdl arguments."
        )
    if backend == "cublaslt" and (enable_bias or pdl):
        pytest.skip(
            "mm_bf16 with cuBLASLt backend does not support bias or pdl arguments."
        )
    if backend == "cutile" and (enable_bias or pdl):
        pytest.skip(
            "mm_bf16 with cuTile backend does not support bias or pdl arguments."
        )
    if res_dtype != torch.bfloat16 and backend == "tgv":
        pytest.skip(
            "mm_bf16 with TGV backend does not support specifying non-bfloat16 result dtypes."
        )
    if res_dtype != torch.bfloat16 and backend == "tinygemm":
        pytest.skip(
            "mm_bf16 with TinyGEMM backend does not support specifying non-bfloat16 result dtypes."
        )
    # cuDNN on SM103 does not support bf16 input -> fp16 output
    if (
        backend == "cudnn"
        and compute_capability_number == 103
        and res_dtype == torch.float16
    ):
        pytest.skip("cuDNN bf16 GEMM with fp16 output not supported on SM103.")

    torch.manual_seed(42)
    input = torch.randn([m, k], device="cuda", dtype=torch.bfloat16)
    mat2 = torch.randn([n, k], device="cuda", dtype=torch.bfloat16)

    if enable_bias:
        bias = torch.randn(n, device="cuda", dtype=torch.bfloat16)
        reference = F.linear(input, mat2, bias)
    else:
        bias = None
        reference = torch.mm(input, mat2.T)

    out = torch.empty([m, n], device="cuda", dtype=res_dtype)
    with autotune(auto_tuning):
        mm_bf16(input, mat2.T, bias, pdl, out, res_dtype, backend)

    cos_sim = F.cosine_similarity(reference.reshape(-1), out.reshape(-1), dim=0)
    assert cos_sim > 0.99


def test_mm_bf16_cutile_rejects_bias_and_pdl():
    """The v1 cuTile path is alpha=1/beta=0 and ignores bias / pdl — must raise.

    Output dtype, on the other hand, supports bf16 / fp16 / fp32 via the
    polymorphic store epilogue, so it is exercised in the main parametrized
    matrix above rather than rejected here.
    """
    compute_capability = get_compute_capability(torch.device("cuda"))
    cc_num = compute_capability[0] * 10 + compute_capability[1]
    if not mm_bf16.is_backend_supported("cutile", cc_num):
        pytest.skip("cuTile backend not supported on current compute capability.")
    if not is_cuda_tile_available():
        pytest.skip("cuda-tile / tileiras compiler not available in this environment.")

    # a is (m, k) = (64, 1024); b is (n, k) = (2048, 1024); b.T is (k, n) = (1024, 2048)
    a = torch.randn(64, 1024, device="cuda", dtype=torch.bfloat16)
    b = torch.randn(2048, 1024, device="cuda", dtype=torch.bfloat16)
    bias = torch.randn(2048, device="cuda", dtype=torch.bfloat16)
    out = torch.empty(64, 2048, device="cuda", dtype=torch.bfloat16)
    with pytest.raises(ValueError, match="ignores `bias`"):
        mm_bf16(a, b.T, bias, False, out, torch.bfloat16, backend="cutile")
    with pytest.raises(ValueError, match="ignores `pdl`"):
        mm_bf16(a, b.T, None, True, out, torch.bfloat16, backend="cutile")


def test_mm_bf16_cutile_repeat_uses_tune_cache():
    """Second call on the same shape must reuse the cached autotune result.

    Cache hits reproduce the first call's output bit-for-bit (no kernel
    re-tune, no autotune RNG state to diverge). We use ``assert_close`` with
    rtol=atol=0 — a strict-inequality threshold like ``cos_sim > 0.999`` is
    unsafe in bf16 because the rhs rounds to 1.0 and ``1.0 > 1.0`` is False
    even when the kernel is correct.
    """
    compute_capability = get_compute_capability(torch.device("cuda"))
    cc_num = compute_capability[0] * 10 + compute_capability[1]
    if not mm_bf16.is_backend_supported("cutile", cc_num):
        pytest.skip("cuTile backend not supported on current compute capability.")
    if not is_cuda_tile_available():
        pytest.skip("cuda-tile / tileiras compiler not available in this environment.")

    # Shape (m, n, k) = (64, 2048, 1024). mm_bf16(a, b.T) → (m, n).
    a = torch.randn(64, 1024, device="cuda", dtype=torch.bfloat16)
    b = torch.randn(2048, 1024, device="cuda", dtype=torch.bfloat16)
    out = torch.empty(64, 2048, device="cuda", dtype=torch.bfloat16)

    # First call: warms tune cache.
    mm_bf16(a, b.T, None, False, out, torch.bfloat16, backend="cutile")
    out_first = out.clone()
    # Second call: must produce the same result and not raise.
    mm_bf16(a, b.T, None, False, out, torch.bfloat16, backend="cutile")
    torch.testing.assert_close(
        out_first,
        out,
        rtol=0,
        atol=0,
        msg="tune-cache reuse produced divergent output",
    )


def test_cublaslt_bf16_runner_zero_algos():
    """CublasltBf16GemmRunner.forward() must raise when heuristic returns 0 algorithms."""
    from flashinfer.gemm.gemm_base import get_mm_bf16_cublaslt_module
    from flashinfer.utils import get_compute_capability

    compute_capability = get_compute_capability(torch.device("cuda"))
    cc_num = compute_capability[0] * 10 + compute_capability[1]
    if not mm_bf16.is_backend_supported("cublaslt", cc_num):
        pytest.skip("cublaslt backend not supported on this GPU")

    runner = get_mm_bf16_cublaslt_module().cublaslt_bf16_gemm_runner()

    m, n, k = 16, 1024, 1024
    a = torch.randn(m, k, device="cuda", dtype=torch.bfloat16)
    b = torch.randn(n, k, device="cuda", dtype=torch.bfloat16).transpose(-2, -1)
    out = torch.empty(m, n, device="cuda", dtype=torch.bfloat16)
    workspace = torch.empty(32 * 1024 * 1024, device="cuda", dtype=torch.uint8)
    inputs = [a, b, None, None, out, workspace]

    zero_algo_buf = torch.empty(0, dtype=torch.uint8, device="cpu")
    original_get_algos = runner._get_algos
    runner._get_algos = lambda _inputs: (zero_algo_buf, 0)
    try:
        with pytest.raises(RuntimeError, match="zero algorithms"):
            runner.forward(inputs)
    finally:
        runner._get_algos = original_get_algos


@pytest.mark.parametrize("pdl", [False, True])
def test_mm_bf16_cute_dsl_direct_fallback(pdl: bool):
    """The no-autotune fallback uses direct in its measured crossover band."""
    if get_compute_capability(torch.device("cuda")) not in ((10, 0), (10, 3)):
        pytest.skip("CuTeDSL low-M backend requires SM100/SM103.")

    from flashinfer.cute_dsl.utils import is_cute_dsl_available

    if not is_cute_dsl_available():
        pytest.skip("nvidia-cutlass-dsl is not available.")

    m, n, k = 1, 256, 8192
    a = torch.randn((m, k), device="cuda", dtype=torch.bfloat16)
    b = torch.randn((n, k), device="cuda", dtype=torch.bfloat16)
    out = mm_bf16(a, b.T, pdl=pdl, backend="cute-dsl")
    reference = a @ b.T
    cos_sim = F.cosine_similarity(reference.reshape(-1), out.reshape(-1), dim=0)
    assert cos_sim > 0.99


@pytest.mark.parametrize("enable_bias", [False, True])
def test_mm_bf16_cute_dsl_one_stage_splitk(enable_bias: bool):
    """K=128 exercises the shallow one-stage tensor-core pipeline."""
    if get_compute_capability(torch.device("cuda")) not in ((10, 0), (10, 3)):
        pytest.skip("CuTeDSL low-M backend requires SM100/SM103.")

    from flashinfer.cute_dsl.utils import is_cute_dsl_available

    if not is_cute_dsl_available():
        pytest.skip("nvidia-cutlass-dsl is not available.")

    m, n, k = 1, 1024, 128
    a = torch.randn((m, k), device="cuda", dtype=torch.bfloat16)
    b = torch.randn((n, k), device="cuda", dtype=torch.bfloat16)
    bias = (
        torch.randn((n,), device="cuda", dtype=torch.bfloat16) if enable_bias else None
    )
    out = mm_bf16(a, b.T, bias=bias, pdl=True, backend="cute-dsl")
    reference = F.linear(a, b, bias)
    cos_sim = F.cosine_similarity(reference.reshape(-1), out.reshape(-1), dim=0)
    assert cos_sim > 0.99


if __name__ == "__main__":
    pytest.main([__file__])
