"""Tests for Fused RMSNorm Kernel.

Tests correctness and performance of the fused Root Mean Square
Layer Normalization operation.
"""

import pytest
import torch

from sagellm_backend.kernels import FusedRMSNormKernel


class TestFusedRMSNorm:
    """Test suite for Fused RMSNorm kernel."""

    def test_kernel_name(self):
        """Test kernel has correct name."""
        kernel = FusedRMSNormKernel()
        assert kernel.name == "fused_rmsnorm"

    def test_output_shape(self):
        """Test output has correct shape."""
        kernel = FusedRMSNormKernel()

        x = torch.randn(4, 128, 512)
        weight = torch.randn(512)

        output = kernel(x, weight)

        assert output.shape == x.shape

    def test_correctness_vs_unfused(self):
        """Test that fused version matches unfused implementation."""
        kernel = FusedRMSNormKernel()

        def unfused_rmsnorm(x, weight, eps=1e-6):
            """Reference unfused RMSNorm implementation."""
            variance = x.pow(2).mean(dim=-1, keepdim=True)
            x_normalized = x * torch.rsqrt(variance + eps)
            return x_normalized * weight

        # Test various shapes
        for shape in [(32, 128), (4, 64, 256), (2, 32, 1024)]:
            x = torch.randn(*shape)
            weight = torch.randn(shape[-1])
            eps = 1e-6

            # Fused version
            fused_output = kernel(x, weight, eps=eps)

            # Unfused version
            unfused_output = unfused_rmsnorm(x, weight, eps=eps)

            # Should match within floating point tolerance
            assert torch.allclose(fused_output, unfused_output, rtol=1e-5, atol=1e-6)

    def test_cpu_fallback(self):
        """Test CPU fallback implementation."""
        kernel = FusedRMSNormKernel()

        x = torch.randn(4, 128, device="cpu")
        weight = torch.randn(128, device="cpu")

        output = kernel(x, weight)

        # Should work on CPU
        assert output.device.type == "cpu"
        assert output.shape == x.shape

    @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
    def test_cuda_execution(self):
        """Test CUDA execution."""
        kernel = FusedRMSNormKernel()

        x = torch.randn(4, 128, 512, device="cuda")
        weight = torch.randn(512, device="cuda")

        output = kernel(x, weight)

        assert output.device.type == "cuda"
        assert output.shape == x.shape

    def test_different_dtypes(self):
        """Test with different data types."""
        kernel = FusedRMSNormKernel()

        for dtype in [torch.float32, torch.float16]:
            if dtype == torch.float16 and not torch.cuda.is_available():
                continue  # FP16 requires CUDA

            device = "cuda" if torch.cuda.is_available() and dtype == torch.float16 else "cpu"

            x = torch.randn(4, 128, dtype=dtype, device=device)
            weight = torch.randn(128, dtype=dtype, device=device)

            output = kernel(x, weight)

            assert output.dtype == dtype
            assert output.shape == x.shape

    def test_different_eps_values(self):
        """Test with different epsilon values."""
        kernel = FusedRMSNormKernel()

        x = torch.randn(4, 128)
        weight = torch.randn(128)

        for eps in [1e-6, 1e-5, 1e-8]:
            output = kernel(x, weight, eps=eps)
            assert output.shape == x.shape

    def test_zero_variance(self):
        """Test with zero variance input (edge case)."""
        kernel = FusedRMSNormKernel()

        # All zeros should remain zeros after normalization
        x = torch.zeros(4, 128)
        weight = torch.randn(128)

        output = kernel(x, weight, eps=1e-6)

        # With eps, output should be finite (not NaN)
        assert torch.isfinite(output).all()

    def test_large_inputs(self):
        """Test with large inputs."""
        kernel = FusedRMSNormKernel()

        # Large tensors
        x = torch.randn(16, 2048, 4096)
        weight = torch.randn(4096)

        output = kernel(x, weight)

        assert output.shape == x.shape

    def test_gradient_flow(self):
        """Test that gradients flow correctly."""
        kernel = FusedRMSNormKernel()

        x = torch.randn(4, 128, requires_grad=True)
        weight = torch.randn(128, requires_grad=True)

        output = kernel(x, weight)
        loss = output.sum()
        loss.backward()

        # Gradients should exist
        assert x.grad is not None
        assert weight.grad is not None

        # Gradients should be non-zero
        assert not torch.allclose(x.grad, torch.zeros_like(x.grad))
        assert not torch.allclose(weight.grad, torch.zeros_like(weight.grad))

    def test_numerical_stability(self):
        """Test numerical stability with extreme values."""
        kernel = FusedRMSNormKernel()

        # Very large values
        x_large = torch.randn(4, 128) * 1000
        weight = torch.randn(128)

        output_large = kernel(x_large, weight)
        assert torch.isfinite(output_large).all()

        # Very small values
        x_small = torch.randn(4, 128) * 1e-5
        output_small = kernel(x_small, weight)
        assert torch.isfinite(output_small).all()

    @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
    def test_performance_benchmark(self):
        """Benchmark fused vs unfused performance."""
        kernel = FusedRMSNormKernel()

        def unfused_rmsnorm(x, weight, eps=1e-6):
            """Reference unfused RMSNorm implementation."""
            variance = x.pow(2).mean(dim=-1, keepdim=True)
            x_normalized = x * torch.rsqrt(variance + eps)
            return x_normalized * weight

        # Use realistic LLM sizes
        batch_size = 32
        seq_len = 1  # Decode
        hidden_size = 4096

        x = torch.randn(batch_size, seq_len, hidden_size, device="cuda")
        weight = torch.randn(hidden_size, device="cuda")
        eps = 1e-6

        # Warmup
        for _ in range(10):
            _ = kernel(x, weight, eps=eps)
            _ = unfused_rmsnorm(x, weight, eps=eps)

        torch.cuda.synchronize()

        # Benchmark fused
        import time

        start = time.perf_counter()
        for _ in range(100):
            _ = kernel(x, weight, eps=eps)
        torch.cuda.synchronize()
        fused_time = time.perf_counter() - start

        # Benchmark unfused
        start = time.perf_counter()
        for _ in range(100):
            _ = unfused_rmsnorm(x, weight, eps=eps)
        torch.cuda.synchronize()
        unfused_time = time.perf_counter() - start

        speedup = unfused_time / fused_time
        print("\nFused RMSNorm Performance:")
        print(f"  Unfused: {unfused_time * 1000:.3f} ms")
        print(f"  Fused:   {fused_time * 1000:.3f} ms")
        print(f"  Speedup: {speedup:.2f}x")

        # Fused should be faster (typically 1.5-3x)
        # Allow some variance in CI environments
        assert speedup >= 0.9

    def test_batch_consistency(self):
        """Test that each batch element is normalized independently."""
        kernel = FusedRMSNormKernel()

        x = torch.randn(4, 128)
        weight = torch.randn(128)

        # Process as batch
        batch_output = kernel(x, weight)

        # Process individually
        individual_outputs = []
        for i in range(4):
            out = kernel(x[i : i + 1], weight)
            individual_outputs.append(out)

        individual_output = torch.cat(individual_outputs, dim=0)

        # Should match
        assert torch.allclose(batch_output, individual_output, rtol=1e-5, atol=1e-6)

    def test_weight_shape_mismatch_error(self):
        """Test that mismatched weight shape raises error."""
        kernel = FusedRMSNormKernel()

        x = torch.randn(4, 128)
        wrong_weight = torch.randn(64)  # Wrong size

        # Should raise error
        with pytest.raises((RuntimeError, ValueError)):
            kernel(x, wrong_weight)


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