import itertools
import math
import unittest

import torch
from utils import (
    BLOCK_K,
    BLOCK_N,
    factor_for_scale,
    fp8_max,
    fp8_min,
    per_token_quant_int8,
    precision,
    scaled_weight,
    torch_naive_moe,
    torch_w8a8_per_column_moe,
)

from sglang.test.test_utils import CustomTestCase

torch.manual_seed(1234)


class TestSharedExpert(CustomTestCase):
    M = [2, 121]
    N = [32, 32 * 4]
    K = [32, 32 * 2]
    routed_scaling_factor = [16]
    apply_scaling_factor = [True, False]

    M_fp8 = [2, 12]
    N_fp8 = [512]
    K_fp8 = [256]

    def _bf16_shared_expert(self, m, n, k, routed_scaling_factor, apply_scaling_factor):
        dtype = torch.bfloat16

        hidden_states = torch.randn(m, k, dtype=dtype) / k
        w1 = torch.randn(2 * n, k, dtype=dtype)
        w2 = torch.randn(k, n, dtype=dtype)
        fused_output = (
            torch.randn(m, k, dtype=dtype) / k if apply_scaling_factor else None
        )
        routed_scaling_factor = routed_scaling_factor if apply_scaling_factor else None

        # fused moe mutates content in hs
        hidden_states2 = hidden_states.clone()

        # bfloat16
        ref = torch_naive_moe(
            hidden_states,
            w1,
            w2,
            fused_output,
            routed_scaling_factor,
            output_dtype=dtype,
        )
        out = torch.ops.sgl_kernel.shared_expert_cpu(
            hidden_states2,
            w1,
            w2,
            fused_output,
            routed_scaling_factor,
            True,
            False,
            False,
            None,
            None,
            None,
            False,
        )

        atol = rtol = precision[ref.dtype]
        torch.testing.assert_close(ref, out, atol=atol, rtol=rtol)

    def test_bf16_shared_expert(self):
        for params in itertools.product(
            self.M,
            self.N,
            self.K,
            self.routed_scaling_factor,
            self.apply_scaling_factor,
        ):
            with self.subTest(
                m=params[0],
                n=params[1],
                k=params[2],
                routed_scaling_factor=params[3],
                apply_scaling_factor=params[4],
            ):
                self._bf16_shared_expert(*params)

    def _int8_shared_expert(self, m, n, k, routed_scaling_factor, apply_scaling_factor):
        dtype = torch.bfloat16

        hidden_states = torch.randn(m, k, dtype=dtype) / k
        w1 = torch.randn(2 * n, k, dtype=dtype)
        w2 = torch.randn(k, n, dtype=dtype)
        fused_output = (
            torch.randn(m, k, dtype=dtype) / k if apply_scaling_factor else None
        )
        routed_scaling_factor = routed_scaling_factor if apply_scaling_factor else None

        # fused moe mutates content in hs
        hidden_states2 = hidden_states.clone()

        w1_q, w1_s = per_token_quant_int8(w1)
        w2_q, w2_s = per_token_quant_int8(w2)
        ref = torch_w8a8_per_column_moe(
            hidden_states,
            w1_q,
            w2_q,
            w1_s,
            w2_s,
            fused_output,
            routed_scaling_factor,
        )
        out = torch.ops.sgl_kernel.shared_expert_cpu(
            hidden_states2,
            w1_q,
            w2_q,
            fused_output,
            routed_scaling_factor,
            True,
            True,
            False,
            w1_s,
            w2_s,
            None,
            False,
        )

        atol = rtol = precision[ref.dtype]
        torch.testing.assert_close(ref, out, atol=atol, rtol=rtol)

    def test_int8_shared_expert(self):
        for params in itertools.product(
            self.M,
            self.N,
            self.K,
            self.routed_scaling_factor,
            self.apply_scaling_factor,
        ):
            with self.subTest(
                m=params[0],
                n=params[1],
                k=params[2],
                routed_scaling_factor=params[3],
                apply_scaling_factor=params[4],
            ):
                self._int8_shared_expert(*params)

    def _fp8_shared_expert(self, m, n, k, routed_scaling_factor, apply_scaling_factor):
        dtype = torch.bfloat16

        hidden_states = torch.randn(m, k, dtype=dtype) / math.sqrt(k)

        w1_fp32 = torch.randn(1, 2 * n, k)
        w1 = (w1_fp32 * fp8_max).clamp(min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn)

        w2_fp32 = torch.randn(1, k, n)
        w2 = (w2_fp32 * fp8_max).clamp(min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn)

        w1s = torch.randn(1, 2 * n // BLOCK_N, k // BLOCK_K) * factor_for_scale
        w2s = torch.randn(1, k // BLOCK_N, n // BLOCK_K) * factor_for_scale

        w1_scaled = scaled_weight(w1, w1s).view(2 * n, k)
        w2_scaled = scaled_weight(w2, w2s).view(k, n)

        # change back to 2D
        w1, w2 = w1.squeeze(0), w2.squeeze(0)
        w1s, w2s = w1s.squeeze(0), w2s.squeeze(0)
        w1_scaled, w2_scaled = w1_scaled.squeeze(0), w2_scaled.squeeze(0)

        fused_output = (
            torch.randn(m, k, dtype=dtype) / math.sqrt(k)
            if apply_scaling_factor
            else None
        )
        routed_scaling_factor = routed_scaling_factor if apply_scaling_factor else None
        hidden_states2 = hidden_states.clone()

        # ref with bfloat16
        ref = torch_naive_moe(
            hidden_states,
            w1_scaled,
            w2_scaled,
            fused_output,
            routed_scaling_factor,
            output_dtype=dtype,
        )

        w1 = torch.ops.sgl_kernel.convert_weight_packed(w1)  # [2N, K]
        w2 = torch.ops.sgl_kernel.convert_weight_packed(w2)  # [K, N]
        out = torch.ops.sgl_kernel.shared_expert_cpu(
            hidden_states2,
            w1,
            w2,
            fused_output,
            routed_scaling_factor,
            True,
            False,
            True,
            w1s,
            w2s,
            [BLOCK_N, BLOCK_K],
            True,
        )

        atol = rtol = precision[ref.dtype]
        torch.testing.assert_close(ref, out, atol=atol, rtol=rtol)

    def test_fp8_shared_expert(self):
        for params in itertools.product(
            self.M_fp8,
            self.N_fp8,
            self.K_fp8,
            self.routed_scaling_factor,
            self.apply_scaling_factor,
        ):
            with self.subTest(
                m=params[0],
                n=params[1],
                k=params[2],
                routed_scaling_factor=params[3],
                apply_scaling_factor=params[4],
            ):
                self._fp8_shared_expert(*params)


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