"""Tests for Kunlun XPU kernels.

Tests Kunlun kernel implementations including:
- MatMul operations
- GEMM operations
- Batched operations
"""

from __future__ import annotations

import pytest
import torch

# Skip all tests if torch_xpu is not available
try:
    import torch_xpu  # noqa: F401

    _has_xpu = hasattr(torch, "xpu") and torch.xpu.is_available()
except ImportError:
    _has_xpu = False

pytestmark = pytest.mark.skipif(
    not _has_xpu, reason="Kunlun XPU (torch_xpu) not available. Tests require XPU hardware."
)

from sagellm_backend.kernels.kunlun import KunlunGEMM, KunlunMatMul


class TestKunlunMatMul:
    """Test cases for KunlunMatMul kernel."""

    def test_init(self) -> None:
        """Test MatMul kernel initialization."""
        kernel = KunlunMatMul(device="xpu:0")
        assert kernel is not None

    def test_simple_matmul(self) -> None:
        """Test simple matrix multiplication."""
        kernel = KunlunMatMul(device="xpu:0")

        M, K, N = 4, 8, 6
        input_a = torch.randn(M, K, device="xpu:0")
        input_b = torch.randn(K, N, device="xpu:0")

        output = kernel.forward(input_a, input_b)

        assert output.shape == (M, N)
        assert output.device.type == "xpu"

    def test_matmul_with_transpose(self) -> None:
        """Test matmul with transposed inputs."""
        kernel = KunlunMatMul(device="xpu:0")

        M, K, N = 3, 5, 4
        input_a = torch.randn(K, M, device="xpu:0")  # Transposed
        input_b = torch.randn(N, K, device="xpu:0")  # Transposed

        # A^T @ B^T
        output = kernel.forward(input_a, input_b, transpose_a=True, transpose_b=True)

        assert output.shape == (M, N)

    def test_matmul_with_alpha(self) -> None:
        """Test matmul with scaling factor."""
        kernel = KunlunMatMul(device="xpu:0")

        input_a = torch.randn(2, 3, device="xpu:0")
        input_b = torch.randn(3, 2, device="xpu:0")

        output = kernel.forward(input_a, input_b, alpha=2.0)

        # Result should be 2 * (A @ B)
        expected = 2.0 * torch.matmul(input_a, input_b)
        assert torch.allclose(output, expected, rtol=1e-4, atol=1e-4)

    def test_matmul_with_bias(self) -> None:
        """Test matmul with bias."""
        kernel = KunlunMatMul(device="xpu:0")

        M, K, N = 2, 3, 4
        input_a = torch.randn(M, K, device="xpu:0")
        input_b = torch.randn(K, N, device="xpu:0")
        bias = torch.randn(N, device="xpu:0")

        output = kernel.forward(input_a, input_b, bias=bias)

        assert output.shape == (M, N)

    def test_batched_matmul(self) -> None:
        """Test batched matrix multiplication."""
        kernel = KunlunMatMul(device="xpu:0")

        batch_size, M, K, N = 3, 4, 5, 6
        input_a = torch.randn(batch_size, M, K, device="xpu:0")
        input_b = torch.randn(batch_size, K, N, device="xpu:0")

        output = kernel.forward(input_a, input_b)

        assert output.shape == (batch_size, M, N)

    def test_matmul_cpu_to_xpu(self) -> None:
        """Test automatic device transfer from CPU to XPU."""
        kernel = KunlunMatMul(device="xpu:0")

        input_a = torch.randn(2, 3)  # CPU
        input_b = torch.randn(3, 4)  # CPU

        output = kernel.forward(input_a, input_b)

        assert output.device.type == "xpu"

    def test_matmul_callable(self) -> None:
        """Test that kernel is callable."""
        kernel = KunlunMatMul(device="xpu:0")

        input_a = torch.randn(2, 3, device="xpu:0")
        input_b = torch.randn(3, 4, device="xpu:0")

        # Should be callable directly
        output = kernel(input_a, input_b)
        assert output.shape == (2, 4)


class TestKunlunGEMM:
    """Test cases for KunlunGEMM kernel."""

    def test_init(self) -> None:
        """Test GEMM kernel initialization."""
        kernel = KunlunGEMM(device="xpu:0")
        assert kernel is not None

    def test_simple_gemm(self) -> None:
        """Test simple GEMM operation."""
        kernel = KunlunGEMM(device="xpu:0")

        M, K, N = 4, 8, 6
        input_a = torch.randn(M, K, device="xpu:0")
        input_b = torch.randn(K, N, device="xpu:0")

        output = kernel.forward(input_a, input_b)

        assert output.shape == (M, N)
        assert output.device.type == "xpu"

    def test_gemm_with_alpha_beta(self) -> None:
        """Test GEMM with alpha and beta scaling."""
        kernel = KunlunGEMM(device="xpu:0")

        M, K, N = 3, 4, 5
        input_a = torch.randn(M, K, device="xpu:0")
        input_b = torch.randn(K, N, device="xpu:0")
        output = torch.randn(M, N, device="xpu:0")
        output_orig = output.clone()

        # C = alpha * (A @ B) + beta * C
        result = kernel.forward(input_a, input_b, output, alpha=2.0, beta=0.5)

        # Verify computation
        expected = 2.0 * torch.matmul(input_a, input_b) + 0.5 * output_orig
        assert torch.allclose(result, expected, rtol=1e-4, atol=1e-4)

    def test_gemm_with_transpose(self) -> None:
        """Test GEMM with transposed inputs."""
        kernel = KunlunGEMM(device="xpu:0")

        M, K, N = 3, 5, 4
        input_a = torch.randn(K, M, device="xpu:0")  # Will be transposed
        input_b = torch.randn(N, K, device="xpu:0")  # Will be transposed

        output = kernel.forward(input_a, input_b, transpose_a=True, transpose_b=True)

        assert output.shape == (M, N)

    def test_gemm_beta_zero(self) -> None:
        """Test GEMM with beta=0 (ignore existing output)."""
        kernel = KunlunGEMM(device="xpu:0")

        input_a = torch.randn(2, 3, device="xpu:0")
        input_b = torch.randn(3, 4, device="xpu:0")
        output = torch.ones(2, 4, device="xpu:0") * 999.0  # Should be ignored

        result = kernel.forward(input_a, input_b, output, alpha=1.0, beta=0.0)

        # Should equal A @ B (beta=0 means output is ignored)
        expected = torch.matmul(input_a, input_b)
        assert torch.allclose(result, expected, rtol=1e-4, atol=1e-4)

    def test_batched_gemm(self) -> None:
        """Test batched GEMM operation."""
        kernel = KunlunGEMM(device="xpu:0")

        batch_size, M, K, N = 2, 4, 5, 6
        input_a = torch.randn(batch_size, M, K, device="xpu:0")
        input_b = torch.randn(batch_size, K, N, device="xpu:0")

        output = kernel.batched_gemm(input_a, input_b)

        assert output.shape == (batch_size, M, N)

    def test_gemm_invalid_shapes(self) -> None:
        """Test GEMM with invalid shapes raises error."""
        kernel = KunlunGEMM(device="xpu:0")

        input_a = torch.randn(2, 3, device="xpu:0")
        input_b = torch.randn(5, 4, device="xpu:0")  # Incompatible shape

        with pytest.raises(ValueError, match="Incompatible shapes"):
            kernel.forward(input_a, input_b)

    def test_gemm_callable(self) -> None:
        """Test that kernel is callable."""
        kernel = KunlunGEMM(device="xpu:0")

        input_a = torch.randn(2, 3, device="xpu:0")
        input_b = torch.randn(3, 4, device="xpu:0")

        # Should be callable directly
        output = kernel(input_a, input_b)
        assert output.shape == (2, 4)


@pytest.mark.skipif(not _has_xpu, reason="XPU not available")
def test_matmul_gemm_equivalence() -> None:
    """Test that MatMul and GEMM produce equivalent results."""
    matmul = KunlunMatMul(device="xpu:0")
    gemm = KunlunGEMM(device="xpu:0")

    input_a = torch.randn(4, 5, device="xpu:0")
    input_b = torch.randn(5, 6, device="xpu:0")

    output_matmul = matmul.forward(input_a, input_b)
    output_gemm = gemm.forward(input_a, input_b)

    assert torch.allclose(output_matmul, output_gemm, rtol=1e-4, atol=1e-4)
