from sglang.test.ci.ci_register import register_cuda_ci

register_cuda_ci(est_time=10, suite="stage-b-test-1-gpu-small")

import os
import unittest

import torch

try:
    from sglang.srt.layers.attention.mamba.mamba_state_scatter_triton import (
        fused_mamba_state_scatter_with_mask,
    )

    _FUSED_IMPORT_ERROR = None
except Exception as e:  # pragma: no cover
    fused_mamba_state_scatter_with_mask = None
    _FUSED_IMPORT_ERROR = e


def _dtype_from_str(name: str) -> torch.dtype:
    mapping = {
        "bfloat16": torch.bfloat16,
        "float16": torch.float16,
        "float32": torch.float32,
    }
    if name not in mapping:
        raise ValueError(
            f"Unsupported dtype string {name!r}. Supported: {sorted(mapping.keys())}"
        )
    return mapping[name]


def _ref_scatter(dst, src, dst_indices, src_indices, step_indices):
    """Reference implementation using PyTorch advanced indexing."""
    # dst: [L, C, E]
    # src: [L, S, D, E]
    dst[:, dst_indices] = src[:, src_indices, step_indices].to(dst.dtype, copy=False)


def _ref_update_like(
    ssm_states,
    intermediate_ssm,
    conv_states,
    intermediate_conv,
    *,
    state_indices_tensor,
    accepted_steps,
    mamba_track_indices=None,
    mamba_steps_to_track=None,
):
    """Reference implementation using PyTorch advanced indexing for correctness verification."""
    request_number = accepted_steps.shape[0]
    intermediate_state_indices = torch.arange(
        request_number, dtype=torch.int32, device=accepted_steps.device
    )

    valid_mask = accepted_steps >= 0
    dst_state_indices = state_indices_tensor[valid_mask].to(torch.int64)
    src_state_indices = intermediate_state_indices[valid_mask].to(torch.int64)
    last_steps = accepted_steps[valid_mask].to(torch.int64)

    # Only scatter if there are valid indices (but don't early return -
    # mamba_track_indices processing is independent)
    if dst_state_indices.numel() > 0:
        _ref_scatter(
            ssm_states,
            intermediate_ssm,
            dst_state_indices,
            src_state_indices,
            last_steps,
        )
        _ref_scatter(
            conv_states,
            intermediate_conv,
            dst_state_indices,
            src_state_indices,
            last_steps,
        )

    if mamba_track_indices is not None:
        assert mamba_steps_to_track is not None
        track_mask = mamba_steps_to_track >= 0
        if not track_mask.any():
            return
        dst_track_indices = mamba_track_indices[track_mask].to(torch.int64)
        src_track_indices = intermediate_state_indices[track_mask].to(torch.int64)
        track_steps = mamba_steps_to_track[track_mask].to(torch.int64)

        _ref_scatter(
            ssm_states,
            intermediate_ssm,
            dst_track_indices,
            src_track_indices,
            track_steps,
        )
        _ref_scatter(
            conv_states,
            intermediate_conv,
            dst_track_indices,
            src_track_indices,
            track_steps,
        )


def _fused_update_like(
    ssm_states,
    intermediate_ssm,
    conv_states,
    intermediate_conv,
    *,
    state_indices_tensor,
    accepted_steps,
    mamba_track_indices=None,
    mamba_steps_to_track=None,
):
    """Matches the fully fused logic that avoids index_select and nonzero calls."""
    # Use fully fused kernel that handles masking internally
    fused_mamba_state_scatter_with_mask(
        ssm_states,
        intermediate_ssm,
        state_indices_tensor,
        accepted_steps,
    )
    fused_mamba_state_scatter_with_mask(
        conv_states,
        intermediate_conv,
        state_indices_tensor,
        accepted_steps,
    )

    if mamba_track_indices is not None:
        assert mamba_steps_to_track is not None
        fused_mamba_state_scatter_with_mask(
            ssm_states,
            intermediate_ssm,
            mamba_track_indices,
            mamba_steps_to_track,
        )
        fused_mamba_state_scatter_with_mask(
            conv_states,
            intermediate_conv,
            mamba_track_indices,
            mamba_steps_to_track,
        )


def _time_cuda_ms(fn, iters=50, warmup=10):
    """Measure average CUDA time (ms) using CUDA events."""
    for _ in range(warmup):
        fn()
    torch.cuda.synchronize()

    start = torch.cuda.Event(enable_timing=True)
    end = torch.cuda.Event(enable_timing=True)
    start.record()
    for _ in range(iters):
        fn()
    end.record()
    torch.cuda.synchronize()
    return start.elapsed_time(end) / iters


class TestMambaStateScatterCorrectness(unittest.TestCase):
    @unittest.skipUnless(torch.cuda.is_available(), "CUDA is required for this test.")
    def test_fused_matches_reference(self):
        """Test that fused_mamba_state_scatter_with_mask matches the reference."""
        if fused_mamba_state_scatter_with_mask is None:
            self.skipTest(
                f"fused_mamba_state_scatter_with_mask import failed: {_FUSED_IMPORT_ERROR}"
            )

        torch.manual_seed(42)
        device = torch.device("cuda")

        # Keep sizes moderate so this test is quick.
        L = 8
        B = 32
        C = 49
        D = 5
        ssm_elems = 1024
        conv_elems = 512

        ssm_states0 = torch.randn(
            (L, C, ssm_elems), device=device, dtype=torch.bfloat16
        )
        conv_states0 = torch.randn(
            (L, C, conv_elems), device=device, dtype=torch.bfloat16
        )
        intermediate_ssm = torch.randn(
            (L, B, D, ssm_elems), device=device, dtype=torch.bfloat16
        )
        intermediate_conv = torch.randn(
            (L, B, D, conv_elems), device=device, dtype=torch.bfloat16
        )

        # unique cache lines (no duplicates) to avoid nondeterministic write order
        state_indices_tensor = torch.randperm(C, device=device, dtype=torch.int64)[
            :B
        ].to(torch.int32)

        accepted_steps = torch.randint(0, D, (B,), device=device, dtype=torch.int64)
        # set ~10% invalid
        invalid = torch.rand((B,), device=device) < 0.1
        accepted_steps[invalid] = -1

        # Optional track update
        mamba_track_indices = torch.randperm(C, device=device, dtype=torch.int64)[:B]
        mamba_steps_to_track = torch.randint(
            0, D, (B,), device=device, dtype=torch.int64
        )
        track_invalid = torch.rand((B,), device=device) < 0.7
        mamba_steps_to_track[track_invalid] = -1

        ssm_ref = ssm_states0.clone()
        conv_ref = conv_states0.clone()
        ssm_fused = ssm_states0.clone()
        conv_fused = conv_states0.clone()

        _ref_update_like(
            ssm_ref,
            intermediate_ssm,
            conv_ref,
            intermediate_conv,
            state_indices_tensor=state_indices_tensor,
            accepted_steps=accepted_steps,
            mamba_track_indices=mamba_track_indices,
            mamba_steps_to_track=mamba_steps_to_track,
        )
        _fused_update_like(
            ssm_fused,
            intermediate_ssm,
            conv_fused,
            intermediate_conv,
            state_indices_tensor=state_indices_tensor,
            accepted_steps=accepted_steps,
            mamba_track_indices=mamba_track_indices,
            mamba_steps_to_track=mamba_steps_to_track,
        )

        torch.testing.assert_close(ssm_fused, ssm_ref)
        torch.testing.assert_close(conv_fused, conv_ref)


class TestMambaStateScatterPerf(unittest.TestCase):
    @unittest.skipUnless(torch.cuda.is_available(), "CUDA is required for this test.")
    def test_perf_report_old_vs_fused(self):
        """Optional microbenchmark comparing baseline vs fused kernel.

        Enable with: SGLANG_RUN_MAMBA_SCATTER_PERF_TEST=1
        """
        if os.environ.get("SGLANG_RUN_MAMBA_SCATTER_PERF_TEST", "0") != "1":
            self.skipTest("Set SGLANG_RUN_MAMBA_SCATTER_PERF_TEST=1 to run perf test.")
        if fused_mamba_state_scatter_with_mask is None:
            self.skipTest(
                f"fused_mamba_state_scatter_with_mask import failed: {_FUSED_IMPORT_ERROR}"
            )

        torch.manual_seed(0)
        device = torch.device("cuda")

        # Parameterize sizes via env vars so we can match a real model more closely.
        L = int(os.environ.get("SGLANG_MAMBA_SCATTER_LAYERS", "32"))
        B = int(os.environ.get("SGLANG_MAMBA_SCATTER_BATCH", "48"))
        C = int(os.environ.get("SGLANG_MAMBA_SCATTER_CACHE", "49"))
        D = int(os.environ.get("SGLANG_MAMBA_SCATTER_DRAFT_TOKENS", "5"))
        ssm_elems = int(os.environ.get("SGLANG_MAMBA_SCATTER_SSM_ELEMS", "4096"))
        conv_elems = int(os.environ.get("SGLANG_MAMBA_SCATTER_CONV_ELEMS", "512"))
        invalid_ratio = float(
            os.environ.get("SGLANG_MAMBA_SCATTER_INVALID_RATIO", "0.0")
        )
        track_ratio = float(os.environ.get("SGLANG_MAMBA_SCATTER_TRACK_RATIO", "0.0"))
        ssm_dtype = _dtype_from_str(
            os.environ.get("SGLANG_MAMBA_SCATTER_SSM_DTYPE", "bfloat16")
        )
        conv_dtype = _dtype_from_str(
            os.environ.get("SGLANG_MAMBA_SCATTER_CONV_DTYPE", "bfloat16")
        )

        # Use zeros for dst so each iteration overwrites the same memory.
        ssm_states = torch.zeros((L, C, ssm_elems), device=device, dtype=ssm_dtype)
        conv_states = torch.zeros((L, C, conv_elems), device=device, dtype=conv_dtype)
        intermediate_ssm = torch.randn(
            (L, B, D, ssm_elems), device=device, dtype=ssm_dtype
        )
        intermediate_conv = torch.randn(
            (L, B, D, conv_elems), device=device, dtype=conv_dtype
        )

        state_indices_tensor = torch.randperm(C, device=device, dtype=torch.int64)[
            :B
        ].to(torch.int32)
        accepted_steps = torch.randint(0, D, (B,), device=device, dtype=torch.int64)
        if invalid_ratio > 0:
            invalid = torch.rand((B,), device=device) < invalid_ratio
            accepted_steps[invalid] = -1

        mamba_track_indices = None
        mamba_steps_to_track = None
        if track_ratio > 0:
            mamba_track_indices = torch.randperm(C, device=device, dtype=torch.int64)[
                :B
            ]
            mamba_steps_to_track = torch.randint(
                0, D, (B,), device=device, dtype=torch.int64
            )
            track_invalid = torch.rand((B,), device=device) >= track_ratio
            mamba_steps_to_track[track_invalid] = -1

        def ref_fn():
            _ref_update_like(
                ssm_states,
                intermediate_ssm,
                conv_states,
                intermediate_conv,
                state_indices_tensor=state_indices_tensor,
                accepted_steps=accepted_steps,
                mamba_track_indices=mamba_track_indices,
                mamba_steps_to_track=mamba_steps_to_track,
            )

        def fused_fn():
            _fused_update_like(
                ssm_states,
                intermediate_ssm,
                conv_states,
                intermediate_conv,
                state_indices_tensor=state_indices_tensor,
                accepted_steps=accepted_steps,
                mamba_track_indices=mamba_track_indices,
                mamba_steps_to_track=mamba_steps_to_track,
            )

        # Warm up JIT compilation for triton kernels (and caches for torch indexing)
        ref_fn()
        fused_fn()
        torch.cuda.synchronize()

        ref_ms = _time_cuda_ms(ref_fn)
        fused_ms = _time_cuda_ms(fused_fn)

        num_valid = int((accepted_steps >= 0).sum().item())
        ratio = fused_ms / ref_ms if ref_ms > 0 else float("inf")
        speedup = ref_ms / fused_ms if fused_ms > 0 else float("inf")

        # Print a concise report
        print(
            "\n[MambaStateScatterPerf]\n"
            f"  shapes: L={L} B={B} C={C} D={D} ssm_elems={ssm_elems} conv_elems={conv_elems}\n"
            f"  dtypes: ssm={ssm_dtype} conv={conv_dtype}\n"
            f"  valid: {num_valid}/{B}  invalid_ratio={invalid_ratio}  track_ratio={track_ratio}\n"
            f"  ref_total_ms (baseline):  {ref_ms:.4f}\n"
            f"  fused_total_ms:           {fused_ms:.4f}  (ratio={ratio:.3f}x, speedup={speedup:.2f}x)\n"
        )


if __name__ == "__main__":  # pragma: no cover
    unittest.main()
