import importlib.util
import unittest

try:
    import torch
except ImportError:
    torch = None

GPU = torch is not None and torch.cuda.is_available() and importlib.util.find_spec("vllm") is not None


@unittest.skipUnless(GPU, "needs vLLM and a GPU")
class Kernels(unittest.TestCase):
    def setUp(self):
        torch.manual_seed(0)

    def test_unpack_writes_what_to_gpu_int4_computes(self):
        from prain.kernels import unpack_int4
        from prain.quant import to_gpu_int4

        for n, half in ((256, 128), (1408, 1408), (2816, 352)):
            cpu_layout = torch.randint(0, 256, (n, half), dtype=torch.uint8, device="cuda")
            dst = torch.zeros_like(cpu_layout)
            unpack_int4(cpu_layout.flatten(), dst)
            self.assertTrue(torch.equal(dst, to_gpu_int4(cpu_layout)), (n, half))

    def test_decode_kernel_computes_what_fused_experts_computes(self):
        from vllm.config import VllmConfig, set_current_vllm_config
        from vllm.model_executor.layers.fused_moe.activation import MoEActivation
        from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts

        from prain.kernels import moe_int4_decode
        from prain.quant import dequantize, moe_quant_config, quantize
        from prain.tiered import cpu_experts, slot_table

        hidden, inter, experts, topk = 256, 128, 6, 4
        w13 = torch.randn(experts, 2 * inter, hidden) * 0.1
        w2 = torch.randn(experts, hidden, inter) * 0.1
        parts = [quantize(w, "int4") for w in (w13, w2)]  # gpu layout
        held = [1, 3, 4]  # the round's experts, in slots 0..2
        slots = [[q[held].cuda(), s[held].cuda()] for q, s in parts]
        table = slot_table({e: i for i, e in enumerate(held)}, experts, "cuda")
        for tokens in (1, 2, 3, 4):
            x = torch.randn(tokens, hidden, dtype=torch.bfloat16)
            ids = torch.stack([torch.randperm(experts)[:topk] for _ in range(tokens)]).int()
            ids[0, -1] = experts  # padding route: the table's last row
            w = torch.rand(tokens, topk)
            got = moe_int4_decode(x.cuda(), w.cuda(), ids.cuda(), table, *slots[0], *slots[1])
            with set_current_vllm_config(VllmConfig()):
                vllm = fused_experts(x.cuda(), slots[0][0], slots[1][0], w.cuda(), ids.cuda(),
                                     activation=MoEActivation.GELU_TANH,
                                     global_num_experts=table.numel(), expert_map=table,
                                     quant_config=moe_quant_config("int4", *[s[1] for s in slots]))
            exact = cpu_experts(x, w, torch.where(ids < experts, ids, -1), set(held), lambda e: tuple(
                dequantize(q, s, "int4")[e] for q, s in parts))
            self.assertEqual((got.shape, got.dtype), ((tokens, hidden), torch.bfloat16))
            torch.testing.assert_close(got.float().cpu(), exact, rtol=3e-2, atol=3e-2)
            torch.testing.assert_close(got.float(), vllm.float(), rtol=3e-2, atol=3e-2)

    def test_staged_copy_fills_a_slot_like_the_repack_path(self):
        from prain.experts import _Staging, stage_int4
        from prain.quant import quantize, to_cpu_int4, to_gpu_int4

        w13, w2 = torch.randn(2, 256, 256), torch.randn(2, 256, 128)
        pack = [[to_cpu_int4(q), s] for q, s in (quantize(w, "int4") for w in (w13, w2))]
        slots = [[torch.zeros((3, *t.shape[1:]), dtype=t.dtype, device="cuda") for t in mats]
                 for mats in pack]
        ring = _Staging(sum(t[0].nbytes for mats in pack for t in mats), "cuda")
        stream = torch.cuda.current_stream()
        for expert, slot in ((1, 2), (0, 0), (1, 1), (0, 2), (1, 0)):  # wraps the ring
            stage_int4([[t[expert] for t in mats] for mats in pack], slots, slot, ring, stream)
            torch.cuda.synchronize()
            for mats, dsts in zip(pack, slots):
                self.assertTrue(torch.equal(dsts[0][slot], to_gpu_int4(mats[0][expert].cuda())))
                self.assertTrue(torch.equal(dsts[1][slot], mats[1][expert].cuda()))


    def test_int4_rounds_on_rdna2_match_the_bf16_kernel_even_when_fp16_overflows(self):
        from types import SimpleNamespace

        from vllm.config import VllmConfig, set_current_vllm_config
        from vllm.model_executor.layers.fused_moe.activation import MoEActivation
        from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts

        from prain.experts import TieredRoutedExperts
        from prain.quant import moe_quant_config, quantize
        from prain.tiered import slot_table

        hidden, inter, experts = 256, 128, 4
        parts = [quantize(w, "int4") for w in (torch.randn(experts, 2 * inter, hidden) * 0.1,
                                                 torch.randn(experts, hidden, inter) * 0.1)]
        slots = [[q.cuda(), s.cuda()] for q, s in parts]
        config = moe_quant_config("int4", slots[0][1], slots[1][1])
        layer = SimpleNamespace(_slots=slots, activation=MoEActivation.GELU_TANH, _qconfig=config)
        table = slot_table({e: e for e in range(experts)}, experts, "cuda")
        with set_current_vllm_config(VllmConfig()):
            for tokens, scale in ((2, 1.0), (64, 1.0), (64, 3000.0)):  # decode kernel, fp16, overflow
                x = (torch.randn(tokens, hidden) * scale).bfloat16().cuda()
                ids = torch.stack([torch.randperm(experts)[:2] for _ in range(tokens)]).int().cuda()
                w = torch.rand(tokens, 2).cuda()
                got = TieredRoutedExperts._int4_gfx10(layer, x, w, ids, table)
                want = fused_experts(x, slots[0][0], slots[1][0], w, ids, activation=layer.activation,
                                     global_num_experts=table.numel(), expert_map=table,
                                     quant_config=config)
                self.assertTrue(torch.isfinite(got).all(), (tokens, scale))
                torch.testing.assert_close(got.float(), want.float(), rtol=3e-2, atol=3e-2 * scale,
                                           msg=f"{tokens} tokens, x * {scale}")


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