"""Tests for DCP LSE combine kernels.

Covers:
1. Triton LSE combine kernel correctness vs CPU reference (base-e and base-2)
2. Various DCP world sizes (N=1,2,4,8)
3. Edge cases: single shard, dominant LSE, equal LSE
4. return_lse mode
5. dcp_a2a_lse_reduce with pre-allocated CUDA graph buffers
"""

import unittest
from unittest.mock import MagicMock

import torch

from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.test_utils import CustomTestCase

register_cuda_ci(est_time=10, stage="base-b", runner_config="1-gpu-large")


class TestLSECombineTritonVsCPU(CustomTestCase):
    """Test Triton LSE combine kernel against CPU reference."""

    @classmethod
    def setUpClass(cls):
        if not torch.cuda.is_available():
            raise unittest.SkipTest("CUDA required for Triton kernel tests")
        cls.device = "cuda"

    def _run_combine_test(
        self, N, B, H_local, D, is_base_e, dtype=torch.bfloat16, atol=1e-2
    ):
        from sglang.kernels.ops.attention.dcp_kernels import (
            _lse_weighted_combine_cpu,
            dcp_lse_combine_triton,
        )

        torch.manual_seed(42)

        partial_outputs = torch.randn(N, B, H_local, D, device=self.device, dtype=dtype)
        if is_base_e:
            partial_lses = torch.randn(
                N, B, H_local, device=self.device, dtype=torch.float32
            )
        else:
            partial_lses = (
                torch.randn(N, B, H_local, device=self.device, dtype=torch.float32)
                * 5.0
            )

        cpu_result = _lse_weighted_combine_cpu(
            partial_outputs.cpu(),
            partial_lses.cpu(),
            is_lse_base_on_e=is_base_e,
        )

        triton_result, _ = dcp_lse_combine_triton(
            partial_outputs,
            partial_lses,
            is_lse_base_on_e=is_base_e,
            return_lse=False,
        )

        torch.testing.assert_close(
            triton_result.float().cpu(),
            cpu_result.float(),
            atol=atol,
            rtol=1e-2,
        )

    def test_n2_base_e(self):
        self._run_combine_test(N=2, B=4, H_local=8, D=64, is_base_e=True)

    def test_n2_base_2(self):
        self._run_combine_test(N=2, B=4, H_local=8, D=64, is_base_e=False)

    def test_n4_base_e(self):
        self._run_combine_test(N=4, B=8, H_local=16, D=128, is_base_e=True)

    def test_n4_base_2(self):
        self._run_combine_test(N=4, B=8, H_local=16, D=128, is_base_e=False)

    def test_n8_base_e(self):
        self._run_combine_test(N=8, B=4, H_local=8, D=128, is_base_e=True)

    def test_n8_base_2(self):
        self._run_combine_test(N=8, B=4, H_local=8, D=512, is_base_e=False)

    def test_n2_large_batch(self):
        self._run_combine_test(N=2, B=64, H_local=16, D=128, is_base_e=False)

    def test_n4_large_head_dim(self):
        self._run_combine_test(N=4, B=8, H_local=8, D=512, is_base_e=True)

    def test_flashmla_natural_log_lse_correction(self):
        """Natural-log LSEs log(2), log(8) give an 8/10 local weight."""
        from sglang.kernels.ops.attention.dcp_kernels import correct_attn_out

        local_output = torch.tensor(
            [[[10.0]]], device=self.device, dtype=torch.bfloat16
        )
        lses = torch.log(torch.tensor([[[2.0]], [[8.0]]], device=self.device))
        corrected, _ = correct_attn_out(
            local_output,
            lses,
            cp_rank=1,
            ctx=None,
            new_output=torch.empty((1, 1, 1), device=self.device),
            is_lse_base_on_e=True,
        )

        torch.testing.assert_close(
            corrected.cpu(), torch.tensor([[[8.0]]]), atol=1e-5, rtol=1e-5
        )


class TestLSECombineSingleShard(CustomTestCase):
    """N=1 should return input unchanged."""

    @classmethod
    def setUpClass(cls):
        if not torch.cuda.is_available():
            raise unittest.SkipTest("CUDA required")
        cls.device = "cuda"

    def test_single_shard(self):
        from sglang.kernels.ops.attention.dcp_kernels import dcp_lse_combine_triton

        N, B, H_local, D = 1, 4, 8, 64
        partial_outputs = torch.randn(
            N, B, H_local, D, device=self.device, dtype=torch.bfloat16
        )
        partial_lses = torch.randn(
            N, B, H_local, device=self.device, dtype=torch.float32
        )

        triton_result, _ = dcp_lse_combine_triton(
            partial_outputs, partial_lses, is_lse_base_on_e=True
        )

        torch.testing.assert_close(
            triton_result.float().cpu(),
            partial_outputs.squeeze(0).float().cpu(),
            atol=1e-3,
            rtol=1e-3,
        )


class TestLSECombineReturnLSE(CustomTestCase):
    """Verify return_lse=True produces valid global LSE."""

    @classmethod
    def setUpClass(cls):
        if not torch.cuda.is_available():
            raise unittest.SkipTest("CUDA required")
        cls.device = "cuda"

    def test_return_lse(self):
        from sglang.kernels.ops.attention.dcp_kernels import dcp_lse_combine_triton

        N, B, H_local, D = 2, 4, 8, 64
        partial_outputs = torch.randn(
            N, B, H_local, D, device=self.device, dtype=torch.bfloat16
        )
        partial_lses = torch.randn(
            N, B, H_local, device=self.device, dtype=torch.float32
        )

        triton_result, triton_lse = dcp_lse_combine_triton(
            partial_outputs, partial_lses, is_lse_base_on_e=True, return_lse=True
        )

        self.assertIsNotNone(triton_lse)
        self.assertEqual(triton_lse.shape, (B, H_local))
        self.assertFalse(torch.isnan(triton_lse).any())


class TestLSECombineEdgeCases(CustomTestCase):
    """Test edge cases for LSE combine."""

    @classmethod
    def setUpClass(cls):
        if not torch.cuda.is_available():
            raise unittest.SkipTest("CUDA required")
        cls.device = "cuda"

    def test_one_shard_dominant(self):
        """One shard has much larger LSE -- output should be close to that shard."""
        from sglang.kernels.ops.attention.dcp_kernels import (
            _lse_weighted_combine_cpu,
            dcp_lse_combine_triton,
        )

        N, B, H_local, D = 2, 1, 1, 64
        partial_outputs = torch.randn(
            N, B, H_local, D, device=self.device, dtype=torch.bfloat16
        )
        partial_lses = torch.tensor(
            [[[100.0]], [[-100.0]]], device=self.device, dtype=torch.float32
        )

        triton_result, _ = dcp_lse_combine_triton(
            partial_outputs, partial_lses, is_lse_base_on_e=True
        )
        cpu_result = _lse_weighted_combine_cpu(
            partial_outputs.cpu(), partial_lses.cpu(), is_lse_base_on_e=True
        )

        torch.testing.assert_close(
            triton_result.float().cpu(), cpu_result.float(), atol=1e-2, rtol=1e-2
        )
        torch.testing.assert_close(
            triton_result.float().cpu(),
            partial_outputs[0].float().cpu(),
            atol=1e-2,
            rtol=1e-2,
        )

    def test_equal_lse(self):
        """Equal LSE across shards -- output should be mean of outputs."""
        from sglang.kernels.ops.attention.dcp_kernels import dcp_lse_combine_triton

        N, B, H_local, D = 2, 1, 1, 64
        partial_outputs = torch.randn(
            N, B, H_local, D, device=self.device, dtype=torch.bfloat16
        )
        partial_lses = torch.tensor(
            [[[5.0]], [[5.0]]], device=self.device, dtype=torch.float32
        )

        triton_result, _ = dcp_lse_combine_triton(
            partial_outputs, partial_lses, is_lse_base_on_e=True
        )
        expected = partial_outputs.float().mean(dim=0)

        torch.testing.assert_close(
            triton_result.float().cpu(), expected.cpu(), atol=1e-2, rtol=1e-2
        )


class TestLSEBaseByBackend(CustomTestCase):
    """Which attention backends report LSE in natural log."""

    def test_natural_log_lse_backends(self):
        from sglang.srt.models.deepseek_common.attention_forward_methods.forward_mla import (
            is_mla_dcp_lse_base_on_e,
        )

        self.assertTrue(is_mla_dcp_lse_base_on_e("flashmla"))
        self.assertTrue(is_mla_dcp_lse_base_on_e("cutedsl_mla"))
        self.assertFalse(is_mla_dcp_lse_base_on_e("flashinfer_mla"))
        self.assertFalse(is_mla_dcp_lse_base_on_e("tokenspeed_mla"))
        self.assertFalse(is_mla_dcp_lse_base_on_e("trtllm_mla"))
        self.assertFalse(is_mla_dcp_lse_base_on_e(None))


class TestDCPA2AReduceWithCUDAGraphBuffers(CustomTestCase):
    """Test dcp_a2a_lse_reduce with pre-allocated CUDA graph buffers."""

    @classmethod
    def setUpClass(cls):
        if not torch.cuda.is_available():
            raise unittest.SkipTest("CUDA required")
        cls.device = "cuda"

    def _make_mock_group(self, world_size):
        group = MagicMock()
        group.world_size = world_size

        def identity_a2a(output, input_):
            output.copy_(input_)

        group.all_to_all_single = MagicMock(side_effect=identity_a2a)
        return group

    def _make_cuda_graph_buffers(self, N, max_bs, H_per_rank, D, lpd=2):
        """Create fused CUDA graph buffers matching dcp_a2a_lse_reduce API."""
        return {
            "send_combined": torch.empty(
                N, max_bs, H_per_rank, D + lpd, dtype=torch.bfloat16, device=self.device
            ),
            "recv_combined": torch.empty(
                N, max_bs, H_per_rank, D + lpd, dtype=torch.bfloat16, device=self.device
            ),
            "send_lse": torch.empty(
                N, max_bs, H_per_rank, dtype=torch.float32, device=self.device
            ),
            "recv_lse": torch.empty(
                N, max_bs, H_per_rank, dtype=torch.float32, device=self.device
            ),
        }

    def test_cuda_graph_buffers_same_as_dynamic(self):
        from sglang.srt.layers.dcp import dcp_a2a_lse_reduce

        torch.manual_seed(123)
        N, B, H_per_rank, D = 2, 4, 8, 128
        H = H_per_rank * N
        max_bs = 16

        group = self._make_mock_group(N)

        attn_out = torch.randn(B, H, D, device=self.device, dtype=torch.bfloat16)
        attn_lse = torch.randn(B, H, device=self.device, dtype=torch.float32)

        result_dynamic = dcp_a2a_lse_reduce(
            attn_out.clone(), attn_lse.clone(), group, is_lse_base_on_e=True
        )

        cuda_graph_buffers = self._make_cuda_graph_buffers(N, max_bs, H_per_rank, D)

        result_graph = dcp_a2a_lse_reduce(
            attn_out.clone(),
            attn_lse.clone(),
            group,
            is_lse_base_on_e=True,
            cuda_graph_buffers=cuda_graph_buffers,
        )

        torch.testing.assert_close(
            result_graph.float().cpu(),
            result_dynamic.float().cpu(),
            atol=1e-5,
            rtol=1e-5,
        )

    def test_cuda_graph_buffers_partial_batch(self):
        """Buffer max_bs > actual B -- should correctly slice."""
        from sglang.srt.layers.dcp import dcp_a2a_lse_reduce

        torch.manual_seed(789)
        N, B, H_per_rank, D = 2, 3, 8, 128
        H = H_per_rank * N
        max_bs = 32

        group = self._make_mock_group(N)

        attn_out = torch.randn(B, H, D, device=self.device, dtype=torch.bfloat16)
        attn_lse = torch.randn(B, H, device=self.device, dtype=torch.float32)

        cuda_graph_buffers = self._make_cuda_graph_buffers(N, max_bs, H_per_rank, D)

        result = dcp_a2a_lse_reduce(
            attn_out,
            attn_lse,
            group,
            is_lse_base_on_e=True,
            cuda_graph_buffers=cuda_graph_buffers,
        )

        self.assertEqual(result.shape, (B, H_per_rank, D))
        self.assertFalse(torch.isnan(result).any())

    def test_pack_matches_the_copy_formulation_it_replaces(self):
        from sglang.kernels.ops.attention.dcp_kernels import (
            _lse_pack_dim,
            dcp_pack_a2a_send,
        )

        for N, B, H_per_rank, D, dtype in (
            (2, 4, 8, 128, torch.bfloat16),
            (4, 3, 2, 64, torch.float16),
            (8, 1, 12, 512, torch.bfloat16),
        ):
            with self.subTest(N=N, B=B, H_per_rank=H_per_rank, D=D, dtype=dtype):
                H = H_per_rank * N
                lpd = _lse_pack_dim(dtype)
                max_bs = B + 5
                out = torch.randn(B, H, D, device=self.device, dtype=dtype)
                lse = torch.randn(B, H, device=self.device, dtype=torch.float32)

                got = torch.zeros(
                    N, max_bs, H_per_rank, D + lpd, dtype=dtype, device=self.device
                )
                dcp_pack_a2a_send(
                    out,
                    lse,
                    got[:, :, :, :D],
                    got.view(torch.float32)[:, :, :, D // lpd],
                )

                want = torch.zeros_like(got)
                want[:, :B, :, :D] = out.view(B, N, H_per_rank, D).permute(1, 0, 2, 3)
                want[:, :B, :, D:] = (
                    lse.view(B, N, H_per_rank)
                    .permute(1, 0, 2)
                    .contiguous()
                    .view(dtype)
                    .view(N, B, H_per_rank, lpd)
                )
                self.assertTrue(
                    torch.equal(got.view(torch.uint8), want.view(torch.uint8))
                )

                lane = got.view(torch.float32)[:, :B, :, D // lpd]
                self.assertTrue(
                    torch.equal(lane, lse.view(B, N, H_per_rank).permute(1, 0, 2))
                )

    def test_pack_serves_the_split_peer_inside_layout(self):
        from sglang.kernels.ops.attention.dcp_kernels import dcp_pack_a2a_send

        for N, B, H_per_rank, D in ((2, 4, 8, 128), (4, 1, 16, 512)):
            with self.subTest(N=N, B=B, H_per_rank=H_per_rank, D=D):
                H = H_per_rank * N
                out = torch.randn(B, H, D, device=self.device, dtype=torch.bfloat16)
                lse = torch.randn(B, H, device=self.device, dtype=torch.float32)

                partial_o = torch.empty(
                    B, H_per_rank, N, D, dtype=torch.bfloat16, device=self.device
                )
                stats = torch.zeros(
                    B, H_per_rank, N, 2, dtype=torch.float32, device=self.device
                )
                dcp_pack_a2a_send(
                    out,
                    lse,
                    partial_o.permute(2, 0, 1, 3),
                    stats[..., 0].permute(2, 0, 1),
                )

                want_o = out.view(B, N, H_per_rank, D).permute(0, 2, 1, 3)
                want_lse = lse.view(B, N, H_per_rank).permute(0, 2, 1)
                self.assertTrue(
                    torch.equal(
                        partial_o.view(torch.uint8),
                        want_o.contiguous().view(torch.uint8),
                    )
                )
                self.assertTrue(torch.equal(stats[..., 0], want_lse))
                self.assertTrue(
                    torch.equal(stats[..., 1], torch.zeros_like(stats[..., 1]))
                )

    def test_buffers_have_fixed_data_ptrs(self):
        """Pre-allocated buffer data_ptr must not change -- required for graph replay."""
        from sglang.srt.layers.dcp import dcp_a2a_lse_reduce

        N, B, H_per_rank, D = 2, 4, 8, 64
        H = H_per_rank * N
        max_bs = 16

        group = self._make_mock_group(N)

        buffers = self._make_cuda_graph_buffers(N, max_bs, H_per_rank, D)
        send_ptr = buffers["send_combined"].data_ptr()
        recv_ptr = buffers["recv_combined"].data_ptr()

        attn_out = torch.randn(B, H, D, device=self.device, dtype=torch.bfloat16)
        attn_lse = torch.randn(B, H, device=self.device, dtype=torch.float32)

        dcp_a2a_lse_reduce(
            attn_out,
            attn_lse,
            group,
            is_lse_base_on_e=True,
            cuda_graph_buffers=buffers,
        )

        self.assertEqual(buffers["send_combined"].data_ptr(), send_ptr)
        self.assertEqual(buffers["recv_combined"].data_ptr(), recv_ptr)


if __name__ == "__main__":
    unittest.main()