# SPDX-License-Identifier: Apache-2.0
"""Tests for the MOE layers.

Run `pytest tests/kernels/test_moe.py`.
"""

from aiter.ops.shuffle import moe_layout_shuffle_gemm1,moe_layout_shuffle_gemm2
# import _custom_ops as ops
import aiter

import pytest
import torch
from typing import List, Mapping, Optional, Tuple

# from utils import torch_moe
from aiter.fused_moe_c import fused_moe
from scalar_type import ScalarType, scalar_types

import torch.nn.functional as F  # 确保导入 functional

def torch_moe(a, w1, w2, score, topk, expert_map):
    B, D = a.shape
    a = a.view(B, -1, D).repeat(1, topk, 1).reshape(-1, D)
    out = torch.zeros(B * topk, w2.shape[1], dtype=a.dtype, device=a.device)
    score = torch.softmax(score, dim=-1, dtype=torch.float32)
    topk_weight, topk_ids = torch.topk(score, topk)
    topk_weight = topk_weight.view(-1)
    topk_ids = topk_ids.view(-1)
    if expert_map is not None:
        topk_ids = expert_map[topk_ids]
    
    for i in range(w1.shape[0]):
        mask = topk_ids == i
        if mask.sum():
            # 1. 计算 w1 的线性变换
            intermediate = a[mask] @ w1[i].transpose(0, 1)
            
            # 2. 兼容低版本 PyTorch 的 SiLU 实现
            d = intermediate.shape[-1] // 2
            # 方法1：用 F.silu（PyTorch 1.6+ 支持，比 torch.silu 兼容性更广）
            silu_part = F.silu(intermediate[..., :d])  # 优先用 functional 中的 silu
            # 方法2：如果 F.silu 也不存在，手动实现 SiLU: x * sigmoid(x)
            # silu_part = intermediate[..., :d] * torch.sigmoid(intermediate[..., :d])
            
            # 3. 与后半部分相乘
            mul_part = intermediate[..., d:]
            silu_and_mul_result = silu_part * mul_part
            
            # 4. 与 w2 相乘
            out[mask] = silu_and_mul_result @ w2[i].transpose(0, 1)
    
    return (out.view(B, -1, w2.shape[1]) *
            topk_weight.view(B, -1, 1).to(out.dtype)).sum(dim=1)

def quantize_weights(w: torch.Tensor,
                     quant_type: ScalarType,
                     group_size: Optional[int],
                     zero_points: bool = False,
                     ref_zero_points_after_scales: bool = False):
    assert quant_type.is_integer(), \
        "Floating point quantization may work but has not been tested"
    assert not zero_points or group_size is not None, \
        "to have group zero points, group_size must be provided "\
        "(-1 group_size is channelwise)"

    orig_device = w.device
    orig_type = w.dtype
    size_k, size_n = w.shape

    assert w.is_floating_point(), "w must be float"

    if group_size == -1:
        group_size = size_k

    # Reshape to [groupsize, -1]
    if group_size is not None and group_size < size_k:
        w = w.reshape((-1, group_size, size_n))
        w = w.permute(1, 0, 2)
        w = w.reshape((group_size, -1))

    # Compute scale for each group
    max_val = torch.max(w, 0, keepdim=True).values
    min_val = torch.min(w, 0, keepdim=True).values

    max_q_val = quant_type.max()
    min_q_val = quant_type.min()

    w_s = torch.Tensor([1.0]).to(w.device)  # unscaled case
    maybe_w_zp = None
    if group_size is not None:
        if zero_points:
            assert not quant_type.is_signed() and quant_type.max() > 0
            w_s = (max_val - min_val).clamp(min=1e-5) / quant_type.max()
            maybe_w_zp = torch.round(torch.abs(min_val / w_s)) \
                .clamp(min_q_val, max_q_val).int()
        else:
            # If the bias is such that there are no possible negative/positive
            #  values, set the max value to inf to avoid divide by 0
            w_s = torch.max(
                abs(max_val / (max_q_val if max_q_val != 0 else torch.inf)),
                abs(min_val / (min_q_val if min_q_val != 0 else torch.inf)))

    # Quantize
    w_q = torch.round(w / w_s).int() + (maybe_w_zp if zero_points else 0)
    w_q = torch.clamp(w_q, min_q_val, max_q_val)

    # Compute ref (dequantized)
    # For some kernels (namely Machete) the zero-points are applied after the
    # scales are applied, for this case computing the reference in similar way
    # allows us to use tighter error tolerances in our unit tests.
    if ref_zero_points_after_scales and maybe_w_zp is not None:
        w_ref = w_q.to(orig_type) * w_s - maybe_w_zp.to(orig_type) * w_s
    else:
        w_ref = (w_q - (maybe_w_zp if zero_points else 0)).to(orig_type) * w_s

    if quant_type.has_bias():
        w_q += quant_type.bias

    # Restore original shapes
    if group_size is not None and group_size < size_k:

        def reshape_w(w):
            w = w.reshape((group_size, -1, size_n))
            w = w.permute(1, 0, 2)
            w = w.reshape((size_k, size_n)).contiguous()
            return w

        w_q = reshape_w(w_q)
        w_ref = reshape_w(w_ref)
        w_s = w_s.reshape((-1, size_n)).contiguous()

    if maybe_w_zp is not None:
        maybe_w_zp = maybe_w_zp.reshape((-1, size_n)).contiguous()
        maybe_w_zp = maybe_w_zp.to(device=orig_device)

    return (
        w_ref.to(device=orig_device),
        w_q.to(device=orig_device),
        w_s if group_size is not None else None,
        maybe_w_zp,
    )

@pytest.mark.parametrize("m", [1])
# @pytest.mark.parametrize("m", [1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,32,64,128,2048,4096,8192,16384,32768])
# @pytest.mark.parametrize("m", [1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,32,64,128])
# @pytest.mark.parametrize("m", [512,1024,2048,4096,8192,16384,32768])
@pytest.mark.parametrize("n", [256])
@pytest.mark.parametrize("k", [7168])
@pytest.mark.parametrize("e", [256])
@pytest.mark.parametrize("topk", [8])
@pytest.mark.parametrize("ep_size", [1])
@pytest.mark.parametrize("dtype", [torch.float16])
@pytest.mark.parametrize("group_size", [64])
@pytest.mark.parametrize("has_zp", [True])
@pytest.mark.parametrize("weight_bits", [4])
def test_fused_moe_wn16(m: int, n: int, k: int, e: int, topk: int,
                        ep_size: int, dtype: torch.dtype, group_size: int,
                        has_zp: bool, weight_bits: int):
    rand_initialize = True
    print(m, n, k, e, topk, dtype, group_size, has_zp, weight_bits)
    if rand_initialize:
        a = torch.randn((m, k), device="cuda", dtype=dtype) /10
        # a = torch.arange(1, m*k + 1, device="cuda", dtype=dtype).reshape(m, k) /100
        w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) /10
        w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) /10
        score = torch.randn((m, e), device="cuda", dtype=dtype)
        # score[0, 0] = 10000
    else:
        a = torch.ones((m, k), device="cuda", dtype=dtype) / 10
        w1 = torch.ones((e, 2 * n, k), device="cuda", dtype=dtype) / 10
        w2 = torch.ones((e, k, n), device="cuda", dtype=dtype) / 10
        score = torch.ones((m, e), device="cuda", dtype=dtype)       

    if weight_bits == 4:
        pack_factor = 2
        quant_type = scalar_types.uint4 if has_zp else scalar_types.uint4b8
    elif weight_bits == 8:
        pack_factor = 1
        quant_type = scalar_types.uint8 if has_zp else scalar_types.uint8b128

    w1_ref = w1.clone()
    w2_ref = w2.clone()
    w1_qweight = torch.empty((e, 2 * n, k // pack_factor),
                             device="cuda",
                             dtype=torch.uint8)
    w2_qweight = torch.empty((e, k, n // pack_factor),
                             device="cuda",
                             dtype=torch.uint8)
    w1_scales = torch.empty((e, 2 * n, k // group_size),
                            device="cuda",
                            dtype=dtype)
    w2_scales = torch.empty((e, k, n // group_size),
                            device="cuda",
                            dtype=dtype)
    w1_qzeros = torch.empty((e, 2 * n // pack_factor, k // group_size),
                            device="cuda",
                            dtype=torch.uint8)
    w2_qzeros = torch.empty((e, k // pack_factor, n // group_size),
                            device="cuda",
                            dtype=torch.uint8)

    for i in range(e * 2):
        expert_id = i % e
        if i // e == 0:
            w, w_ref, w_qweight, w_scales, w_qzeros = \
                w1, w1_ref, w1_qweight, w1_scales, w1_qzeros
        else:
            w, w_ref, w_qweight, w_scales, w_qzeros = \
                w2, w2_ref, w2_qweight, w2_scales, w2_qzeros
        weight, qweight, scales, qzeros = quantize_weights(
            w[expert_id].T, quant_type, group_size, has_zp, False)
        weight = weight.T
        qweight = qweight.T.contiguous().to(torch.uint8)
        scales = scales.T
        if has_zp:
            qzeros = qzeros.T.contiguous().to(torch.uint8)
            # qzeros = torch.zeros_like(qzeros)
            # qzeros = (torch.rand_like(qzeros, dtype=torch.float32)).to(qzeros.dtype)
            # qzeros = torch.arange(1, qzeros.numel()+1, device=qzeros.device, dtype=qzeros.dtype).reshape(qzeros.shape) / 1000
        if weight_bits == 4:
            # qweight = torch.arange(1, qweight.numel()+1, device=qweight.device, dtype=qweight.dtype).reshape(qweight.shape)
            qweight = qweight[:, 1::2] * 16 + qweight[:, ::2]
            # qweight = torch.ones_like(qweight) * 17
            # qweight = (torch.rand_like(qweight, dtype=torch.float32)).to(qweight.dtype)
            # qweight = torch.arange(1, qweight.numel()+1, device=qweight.device, dtype=qweight.dtype).reshape(qweight.shape)
            # if i == 256:
            #     print("执行后内容：")
            #     print("执行后形状：", qweight.shape)
            #     print(qweight)
            if has_zp:
                qzeros = qzeros[1::2, :] * 16 + qzeros[::2, :]

        # scales = torch.ones_like(scales)
        # scales = torch.arange(1, scales.numel()+1, device=scales.device, dtype=scales.dtype).reshape(scales.shape) / 10000000
        # scales = (torch.rand_like(scales, dtype=torch.float32)).to(scales.dtype)
        w_ref[expert_id] = weight
        w_qweight[expert_id] = qweight
        w_scales[expert_id] = scales
        if has_zp:
            w_qzeros[expert_id] = qzeros

    if ep_size > 1:
        local_e = e // ep_size
        if rand_initialize:
            e_ids = torch.randint(0,
                                e, (local_e, ),
                                device="cuda",
                                dtype=torch.int32)
        else:
            e_ids = torch.ones((local_e, ),
                                device="cuda",
                                dtype=torch.int32)
        e_map = torch.full((e, ), -1, device="cuda", dtype=torch.int32)
        e_map[e_ids] = torch.arange(local_e, device="cuda", dtype=torch.int32)
        w1_ref = w1_ref[e_ids]
        w2_ref = w2_ref[e_ids]
        w1_qweight = w1_qweight[e_ids]
        w2_qweight = w2_qweight[e_ids]
        w1_scales = w1_scales[e_ids]
        w2_scales = w2_scales[e_ids]
        w1_qzeros = w1_qzeros[e_ids]
        w2_qzeros = w2_qzeros[e_ids]
    else:
        e_map = None

    w1_qweight_uint32 = w1_qweight.view(-1).view(torch.uint32)
    # new_shape = (e, 2 * n, k // 128, 16)  # uint32张量的形状
    new_shape = (e, 2 * n // 16, 16, k // 32, 4)  # uint32张量的形状
    w1_qweight_uint32_reshaped = w1_qweight_uint32.view(new_shape)
    w1_qweight_uint32_transposed = w1_qweight_uint32_reshaped.transpose(2, 3).contiguous()
    new_shape = (e, 2 * n // 16, k // 128, 4, 16, 4)
    w1_new_trans = w1_qweight_uint32_transposed.view(new_shape)
    w1_new = w1_new_trans.transpose(1, 2).contiguous()

    w2_qweight_uint32 = w2_qweight.view(-1).view(torch.uint32)
    # new_shape = (e, 2 * n, k // 128, 16)  # uint32张量的形状
    new_shape = (e, k // 16, 16, n // 32, 4)  # uint32张量的形状
    w2_qweight_uint32_reshaped = w2_qweight_uint32.view(new_shape)
    w2_qweight_uint32_transposed = w2_qweight_uint32_reshaped.transpose(2, 3).contiguous()
    new_shape = (e, k // 16, n // 128, 4, 16, 4)
    w2_new_trans = w2_qweight_uint32_transposed.view(new_shape)
    w2_new = w2_new_trans.transpose(1, 2).contiguous()

    # num_cus = torch.cuda.get_device_properties(torch.cuda.current_device()).multi_processor_count
    # print("******************************")
    # print(num_cus)

    # w1_qweight_uint32 = w1_qweight.view(-1).view(torch.uint32)
    # new_shape = (e, k, n // 128, 16)  # uint32张量的形状
    # w1_qweight_uint32_reshaped = w1_qweight_uint32.view(new_shape)
    # w1_qweight_uint32_transposed = w1_qweight_uint32_reshaped.transpose(1, 2).contiguous()

    # w2_qweight_uint32 = w2_qweight.view(-1).view(torch.uint32)
    # new_shape = (e, k, n // 128, 16)  # uint32张量的形状
    # w2_qweight_uint32_reshaped = w2_qweight_uint32.view(new_shape)
    # w2_qweight_uint32_transposed = w2_qweight_uint32_reshaped.transpose(1, 2).contiguous()

    # chunk_size = 16  # 分块大小
    # w1_uint32_chunked = w1_qweight_uint32_transposed.view(e, -1, chunk_size, 2 * n)
    # w1_uint32_reshaped = w1_uint32_chunked.transpose(0, 1).contiguous()

    w1_qweight_uint32 = w1_qweight.view(-1).view(torch.uint32)
    new_shape = (e, 2 * n, (k // pack_factor) // 4)  # uint32张量的形状
    w1_qweight_uint32_reshaped = w1_qweight_uint32.view(new_shape)
    w1_qweight_uint32_transposed = w1_qweight_uint32_reshaped.transpose(1, 2).contiguous()
    chunk_size = 16  # 分块大小
    w1_uint32_chunked = w1_qweight_uint32_transposed.view(e, -1, chunk_size, 2 * n)
    w1_uint32_reshaped = w1_uint32_chunked.transpose(0, 1).contiguous()

    triton_output = fused_moe(a,
                          w1_qweight,
                          w2_qweight,
                          w1_new,
                          w2_new,
                          score,
                          topk,
                          renormalize=False,
                          use_int4_w4a16=weight_bits == 4,
                          use_int8_w8a16=weight_bits == 8,
                          global_num_experts=e,
                          expert_map=e_map,
                          w1_scale=w1_scales,
                          w2_scale=w2_scales,
                          w1_zp=w1_qzeros if has_zp else None,
                          w2_zp=w2_qzeros if has_zp else None,
                          block_shape=[0, group_size])
    torch_output = torch_moe(a, w1_ref, w2_ref, score, topk, e_map)

    # torch_output = fused_moe(a,
    #                       w1_qweight,
    #                       w2_qweight,
    #                       w1_qweight,
    #                       w2_qweight,
    #                       score,
    #                       topk,
    #                       renormalize=False,
    #                       use_int4_w4a16=weight_bits == 4,
    #                       use_int8_w8a16=weight_bits == 8,
    #                       use_int4_w4a16_base = True,
    #                       global_num_experts=e,
    #                       expert_map=e_map,
    #                       w1_scale=w1_scales,
    #                       w2_scale=w2_scales,
    #                       w1_zp=w1_qzeros if has_zp else None,
    #                       w2_zp=w2_qzeros if has_zp else None,
    #                       block_shape=[0, group_size])

    # count_relative_differences(triton_output, torch_output, rel_threshold=0.1, abs_threshold=2e-2, print_errors=True)

    # 打印形状信息
    print(f"Triton输出形状: {triton_output.shape}")
    print(f"PyTorch输出形状: {torch_output.shape}\n")

    # 计算差异张量和掩码
    diff = torch.abs(triton_output - torch_output)
    atol_threshold = 2e-2  # 与断言保持一致的阈值
    match_mask = diff <= atol_threshold  # 匹配的位置（True）和不匹配的位置（False）

    # 展平所有张量以便切片操作
    triton_flat = triton_output.flatten()
    torch_flat = torch_output.flatten()
    diff_flat = diff.flatten()
    match_mask_flat = match_mask.flatten()

    # 分别收集匹配和不匹配的索引
    match_indices = torch.nonzero(match_mask_flat).squeeze().tolist()
    mismatch_indices = torch.nonzero(~match_mask_flat).squeeze().tolist()

    # 确保索引是列表格式（处理只有一个元素的情况）
    if not isinstance(match_indices, list):
        match_indices = [match_indices] if match_indices is not None else []
    if not isinstance(mismatch_indices, list):
        mismatch_indices = [mismatch_indices] if mismatch_indices is not None else []

    # 打印总体统计
    print(f"===== 总体匹配统计 =====")
    print(f"总元素数: {triton_flat.numel()}")
    print(f"匹配元素数: {len(match_indices)} ({len(match_indices)/triton_flat.numel():.2%})")
    print(f"不匹配元素数: {len(mismatch_indices)} ({len(mismatch_indices)/triton_flat.numel():.2%})")
    print(f"最大差异: {diff.max().item():.6f}")
    print(f"平均差异: {diff.mean().item():.6f}\n")

    # 定义打印指定索引数据的函数
    def print_samples(name, indices, max_samples=3, elements_per_sample=20):
        if not indices:
            print(f"没有{name}的数据")
            return
            
        print(f"===== {name}数据示例 =====")
        # 控制最大展示样本数
        num_samples = min(max_samples, len(indices) // elements_per_sample + 1)
        
        for i in range(num_samples):
            start_idx = i * elements_per_sample
            end_idx = start_idx + elements_per_sample
            sample_indices = indices[start_idx:end_idx]
            
            if not sample_indices:
                break
                
            print(f"\n样本 {i+1}（索引 {sample_indices[0]}-{sample_indices[-1]}）:")
            print(f"Triton: {[triton_flat[idx].item() for idx in sample_indices]}")
            print(f"PyTorch: {[torch_flat[idx].item() for idx in sample_indices]}")
            print(f"差异:    {[diff_flat[idx].item() for idx in sample_indices]}")

    # 打印匹配的数据示例
    print_samples("匹配", match_indices, max_samples=5)

    # 打印不匹配的数据示例
    print_samples("不匹配", mismatch_indices, max_samples=3)

    # 保留断言检查
    torch.testing.assert_close(triton_output, torch_output, atol=atol_threshold, rtol=0)