import pytest
import torch
import itertools

import sys
sys.path.append("/home/users/xiaowei.zhang/new-git/aiter")

from enum import Enum
from aiter import ck_moe,ck_moe_get_solutions
from aiter.fused_moe_ck import run_fused_experts_ck_impl
from aiter import asm_moe_get_solutions
from aiter.fused_moe_asm_wna16 import fused_experts_asm_impl, run_fused_experts_asm_impl
from op_tests.utility.scalar_type import ScalarType, scalar_types
from op_tests.utility.utils import quantize_weights
from aiter.fused_moe import torch_moe, fused_topk, torch_moe_blockscale
from aiter.test_common import checkAllclose, perftest
from aiter.utility.mp_tuner import mp_tuner
from aiter import dtypes
import pandas as pd
from aiter import ActivationType
from aiter import pertoken_quant
from einops import rearrange
from typing import Optional, List  # Add this import at the top
from aiter.ops.shuffle import asm_shuffle_weight_b8


dataTypes = {
    "f32": dtypes.fp32,
    "float32": dtypes.fp32,
    "f16": dtypes.fp16,
    "float16": dtypes.fp16,
    "bf16": dtypes.bf16,
    "bfloat16": dtypes.bf16,
    "fp8": dtypes.fp8,
}

def get_dtype(dtype_str: str):
    if dtype_str is None:
        return None
    if dtype_str.startswith("torch"):
        return getattr(torch, dtype_str.split(".")[1])
    if dtype_str in dataTypes:
        return dataTypes[dtype_str]
    else:
        print(">>> Warning! Invalid dtype", dtype_str, "using default dtype f16")
    return None


class MoeQuantType:
    NO_QUANT = "no_quant"
    INT4_W4A16 = "int4_w4a16"
    INT4_W4A8 = "int4_w4a8"
    INT8_W8A8 = "int8_w8a8_block"
    INT8_W8A8_C = "int8_w8a8_channel"
    F8_W8A8 = "f8_w8a8_block"
    F8_W8A8_C = "f8_w8a8_channel"
    
    ALL_TYPES = [NO_QUANT, INT4_W4A16, INT4_W4A8, INT8_W8A8, INT8_W8A8_C, F8_W8A8, F8_W8A8_C]
    
    @classmethod
    def is_valid(cls, qtype_str: str) -> bool:
        return qtype_str in cls.ALL_TYPES
    
    @classmethod
    def get_default(cls) -> str:
        return cls.NO_QUANT
    
    
def get_QuantType(qtype_str: str) -> str:
    if qtype_str is None or not MoeQuantType.is_valid(qtype_str):
        print(f">>> Warning! Invalid quant_type '{qtype_str}', using default {MoeQuantType.NO_QUANT}")
        return MoeQuantType.get_default()
    return qtype_str



class SolutionType:
    CK = "ck"
    Triton = "triton"
    ASM = "asm"

    ALL_TYPES = [CK, Triton, ASM]

    @classmethod
    def is_valid(cls, sol_type: str) -> bool:
        return sol_type in cls.ALL_TYPES
    
    @classmethod
    def get_default(cls) -> str:
        return cls.CK



NUM_WARMUP=10
NUM_ITERS=40
NUM_WARMUP_FAST=10
NUM_ITERS_FAST=20


class MoeProblem:
    def __init__(self, quant_type, indtype, token, inter_dim, model_dim, expert, topk, q_size_n, q_size_k, mp=1):
        print("quant_type:", quant_type)
        print("indtype:", indtype)
        
        self.quant_type = quant_type
        self.indtype = indtype
        self.token = token
        self.inter_dim = inter_dim
        self.model_dim = model_dim
        self.expert = expert
        self.topk = topk
        self.mp = mp
        
        self.use_int8_w8a16 = False
        self.use_int4_w4a16 = False
        self.use_int8_w8a8_block = False
        self.use_int4_w4a8_block = False
        self.use_int8_w8a8_channel = False
        self.use_f8_w8a8_block = False
        self.use_f8_w8a8_channel = False
        self.per_channel_quant = False
        self.w1_zp = None
        self.w2_zp = None
        self.w1_zp_triton = None
        self.w2_zp_triton = None
        self.w1_scale = None
        self.w2_scale = None
        self.a1_scale = None
        self.a2_scale = None
        self.block_shape_n = q_size_n
        self.block_shape_k = q_size_k
        self.block_m = 16

        self.sol_type = SolutionType.get_default()
        self.sol_id = -1
        self.time_us = float("inf")
        self.check_err_ratio = 0.05
        
        self.rtol = 1e-2
        self.atol = 100
        torch.manual_seed(0)

        self.ck_gtimedf = pd.DataFrame.from_dict(
            {}, orient="index", columns=["gtimems"]
        )
        self.asm_gtimedf = pd.DataFrame.from_dict(
            {}, orient="index", columns=["gtimems"]
        )
        self.blob = torch.ones(128 * 1024 * 1024, dtype=torch.float32, device="cuda")
        self.input = torch.randn((token, model_dim), device="cuda", dtype=indtype) / 10
        score = torch.randn((token, expert), device="cuda", dtype=torch.float32)

        if quant_type == MoeQuantType.NO_QUANT:
            self.w1_ref = torch.randn((expert, 2 * inter_dim, model_dim), device="cuda", dtype=indtype)
            self.w2_ref = torch.randn((expert, model_dim, inter_dim), device="cuda", dtype=indtype)            
            self.w1_qweight = self.w1_ref
            self.w2_qweight = self.w2_ref
            #self.w1_qweight = asm_shuffle_weight_b8(self.w1_qweight, 1)
            #self.w2_qweight = asm_shuffle_weight_b8(self.w2_qweight, 2)
        elif quant_type == MoeQuantType.INT8_W8A8 or quant_type == MoeQuantType.INT8_W8A8_C:
            factor_for_scale = 1e-2
            int8_info = torch.iinfo(torch.int8)
            int8_max, int8_min = int8_info.max, int8_info.min
            self.w1_ref = (torch.rand((expert, 2 * inter_dim, model_dim), dtype=indtype, device="cuda") - 0.5) * 2
            self.w2_ref = (torch.rand((expert, model_dim, inter_dim), dtype=indtype, device="cuda") - 0.5) * 2
            if quant_type == MoeQuantType.INT8_W8A8:
                self.w1_qweight = self.w1_ref.clamp(min=int8_min, max=int8_max).to(torch.int8)
                self.w2_qweight = self.w2_ref.clamp(min=int8_min, max=int8_max).to(torch.int8)
                self.use_int8_w8a8_block = True
                self.block_shape_n = 128
                self.block_shape_k = 128
                
                n_tiles_w1 = (2 * inter_dim + self.block_shape_n - 1) // self.block_shape_n
                n_tiles_w2 = (model_dim + self.block_shape_n - 1) // self.block_shape_n
                k_tiles_w1 = (model_dim + self.block_shape_k - 1) // self.block_shape_k
                k_tiles_w2 = (inter_dim + self.block_shape_k - 1) // self.block_shape_k
                self.w1_scale = (torch.rand((expert, n_tiles_w1, k_tiles_w1), dtype=torch.float32, device="cuda") * factor_for_scale)
                self.w2_scale = (torch.rand((expert, n_tiles_w2, k_tiles_w2), dtype=torch.float32, device="cuda") * factor_for_scale)
                #self.w1_qweight = asm_shuffle_weight_b8(self.w1_qweight, 1)
                #self.w2_qweight = asm_shuffle_weight_b8(self.w2_qweight, 2)
            elif quant_type == MoeQuantType.INT8_W8A8_C:
                self.use_int8_w8a8_channel = True
                self.per_channel_quant = True
                max_vals = torch.abs(self.w1_ref.to(torch.float32)).max(dim=-1, keepdim=True)[0]
                max_vals = max_vals.clamp(min=1e-5)
                self.w1_scale = max_vals / 127.0
                self.w1_qweight = (self.w1_ref / max_vals * 127.0).round().clamp(min=-128, max=127).to(torch.int8)
                max_vals = torch.abs(self.w2_ref.to(torch.float32)).max(dim=-1, keepdim=True)[0]
                max_vals = max_vals.clamp(min=1e-5)
                self.w2_scale = max_vals / 127.0
                self.w2_qweight = (self.w2_ref / max_vals * 127.0).round().clamp(min=-128, max=127).to(torch.int8)
                #self.w1_qweight = asm_shuffle_weight_b8(self.w1_qweight, 1)
                #self.w2_qweight = asm_shuffle_weight_b8(self.w2_qweight, 2)

        elif quant_type == MoeQuantType.F8_W8A8 or quant_type == MoeQuantType.F8_W8A8_C:
            self.w1_ref = torch.randn((expert, 2 * inter_dim, model_dim), dtype=indtype, device="cuda") 
            self.w2_ref = torch.randn((expert, model_dim, inter_dim), dtype=indtype, device="cuda")
            if quant_type == MoeQuantType.F8_W8A8:
                self.use_f8_w8a8_block = True
                self.block_shape_n = 128
                self.block_shape_k = 128

                quant_dtype = dtypes.fp8
                tmp = rearrange(
                    self.w1_ref.view(
                        -1,
                        self.w1_ref.shape[1] // self.block_shape_n,
                        self.block_shape_n,
                        self.w1_ref.shape[2] // self.block_shape_k,
                        self.block_shape_k,
                    ),
                    "e num_blk_n blk_n num_blk_k blk_k -> e num_blk_n num_blk_k (blk_n blk_k)",
                ).contiguous()
                self.w1_qweight, self.w1_scale = pertoken_quant(tmp, quant_dtype=quant_dtype)
                self.w1_qweight = rearrange(
                    self.w1_qweight.view(
                        -1,
                        self.w1_ref.shape[1] // self.block_shape_n,
                        self.w1_ref.shape[2] // self.block_shape_k,
                        self.block_shape_n,
                        self.block_shape_k,
                    ),
                    "e num_blk_n num_blk_k blk_n blk_k -> e (num_blk_n blk_n) (num_blk_k blk_k)",
                ).contiguous()
                self.w1_scale = self.w1_scale.view(expert, self.w1_scale.shape[1], self.w1_scale.shape[2])


                # block quant w2
                tmp = rearrange(
                    self.w2_ref.view(
                        -1,
                        self.w2_ref.shape[1] // self.block_shape_n,
                        self.block_shape_n,
                        self.w2_ref.shape[2]// self.block_shape_k,
                        self.block_shape_k,
                    ),
                    "e num_blk_n blk_n num_blk_k blk_k -> e num_blk_n num_blk_k (blk_n blk_k)",
                ).contiguous()
                self.w2_qweight, self.w2_scale = pertoken_quant(tmp, quant_dtype=quant_dtype)
                self.w2_qweight = rearrange(
                    self.w2_qweight.view(
                        -1,
                        self.w2_ref.shape[1] // self.block_shape_n,
                        self.w2_ref.shape[2] // self.block_shape_k,
                        self.block_shape_n,
                        self.block_shape_k,
                    ),
                    "e num_blk_n num_blk_k blk_n blk_k -> e (num_blk_n blk_n) (num_blk_k blk_k)",
                ).contiguous()
                self.w2_scale = self.w2_scale.view(expert, self.w2_scale.shape[1], self.w2_scale.shape[2])
                #shuffle
                #self.w1_qweight = asm_shuffle_weight_b8(self.w1_qweight, 1)
                #self.w2_qweight = asm_shuffle_weight_b8(self.w2_qweight, 2)
            elif quant_type == MoeQuantType.F8_W8A8_C:
                self.use_f8_w8a8_channel = True
                self.per_channel_quant = True

                self.w1_qweight,self.w1_scale = pertoken_quant(self.w1_ref, quant_dtype=dtypes.fp8)
                self.w2_qweight,self.w2_scale = pertoken_quant(self.w2_ref, quant_dtype=dtypes.fp8)
                #shuffle
                #self.w1_qweight = asm_shuffle_weight_b8(self.w1_qweight, 1)
                #self.w2_qweight = asm_shuffle_weight_b8(self.w2_qweight, 2)

        elif quant_type == MoeQuantType.INT4_W4A8 or quant_type == MoeQuantType.INT4_W4A16:
            if quant_type == MoeQuantType.INT4_W4A8:
                self.use_int4_w4a8_block = True
                self.block_shape_n = 0
                self.block_shape_k = 64
            elif quant_type == MoeQuantType.INT4_W4A16:
                self.use_int4_w4a16 = True

            w1 = torch.randn((expert, 2 * inter_dim, model_dim), device="cuda", dtype=indtype) / 10
            w2 = torch.randn((expert, model_dim, inter_dim), device="cuda", dtype=indtype) / 10
            pack_factor = 2
            group_size = 64
            ep_size = 1
            has_zp = True
            quant_type = scalar_types.uint4

            self.w1_ref = w1.clone()
            self.w2_ref = w2.clone()
            self.w1_qweight = torch.empty((expert, 2 * inter_dim, model_dim // pack_factor),device="cuda",dtype=torch.uint8)
            self.w2_qweight = torch.empty((expert, model_dim, inter_dim // pack_factor),device="cuda",dtype=torch.uint8)
            self.w1_scale = torch.empty((expert, 2 * inter_dim, model_dim // group_size),device="cuda",dtype=indtype)
            self.w2_scale = torch.empty((expert, model_dim, inter_dim // group_size),device="cuda",dtype=indtype)
            #asm
            self.w1_zp = torch.empty((expert, 2 * inter_dim, model_dim // group_size // pack_factor),device="cuda",dtype=torch.uint8)
            self.w2_zp = torch.empty((expert, model_dim, inter_dim // group_size // pack_factor),device="cuda",dtype=torch.uint8)
            #triton
            self.w1_zp_triton = torch.empty((expert, 2 * inter_dim // pack_factor, model_dim // group_size),device="cuda",dtype=torch.uint8)
            self.w2_zp_triton = torch.empty((expert, model_dim // pack_factor, inter_dim // group_size),device="cuda",dtype=torch.uint8)

            for i in range(expert * 2):
                expert_id = i % expert
                if i // expert == 0:
                    w, w_ref, w_qweight, w_scales, w_qzeros, w_qzeros_triton = \
                        w1, self.w1_ref, self.w1_qweight, self.w1_scale, self.w1_zp, self.w1_zp_triton
                else:
                    w, w_ref, w_qweight, w_scales, w_qzeros, w_qzeros_triton = \
                        w2, self.w2_ref, self.w2_qweight, self.w2_scale, self.w2_zp, self.w2_zp_triton
                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)
                qweight = qweight[:, 1::2] * 16 + qweight[:, ::2]   # 偶数列存储低4位，奇数列存储高4位
                if has_zp:
                    #asm qzeros
                    qzeros_asm = qzeros[:, 1::2] * 16 + qzeros[:, ::2]
                    #triton qzeros
                    qzeros_triton = qzeros[1::2, :] * 16 + qzeros[::2, :]

                w_ref[expert_id] = weight
                w_qweight[expert_id] = qweight
                w_scales[expert_id] = scales
                if has_zp:
                    w_qzeros[expert_id] = qzeros_asm
                    w_qzeros_triton[expert_id] = qzeros_triton


            if ep_size > 1:
                local_e = expert // ep_size
                e_ids = torch.randint(0,expert, (local_e, ),device="cuda",dtype=torch.int32)
                e_map = torch.full((expert, ), -1, device="cuda", dtype=torch.int32)
                e_map[e_ids] = torch.arange(local_e, device="cuda", dtype=torch.int32)
                self.w1_ref = self.w1_ref[e_ids]
                self.w2_ref = self.w2_ref[e_ids]
                self.w1_qweight = self.w1_qweight[e_ids]
                self.w2_qweight = self.w2_qweight[e_ids]
                self.w1_scale = self.w1_scale[e_ids]
                self.w2_scale = self.w2_scale[e_ids]
                self.w1_zp = self.w1_zp[e_ids]
                self.w2_zp = self.w2_zp[e_ids]
                self.w1_zp_triton = self.w1_zp_triton[e_ids]
                self.w2_zp_triton = self.w2_zp_triton[e_ids]            
            else:
                e_map = None     

        self.topk_weights, self.topk_ids = fused_topk(self.input, score, topk, True)
        if quant_type != MoeQuantType.INT4_W4A8 or quant_type != MoeQuantType.F8_W8A8 or quant_type != MoeQuantType.INT8_W8A8:
            self.ref = torch_moe(self.input, self.w1_ref, self.w2_ref, self.topk_weights, self.topk_ids,
                                        None, None, None, None, None, ActivationType.Silu).to(self.indtype)
        else:
            #block scale version
            self.ref = torch_moe_blockscale(self.input, 
                                            self.w1_qweight, 
                                            self.w2_qweight, 
                                            self.topk_weights, 
                                            self.topk_ids,
                                            indtype,
                                            [self.block_shape_n, self.block_shape_k],
                                            None,
                                            self.w1_scale,
                                            self.w2_scale,
                                            e_map)

    def ck_time_all_sols(self, fast_mode=0):
        num_warmup = NUM_WARMUP_FAST if fast_mode else NUM_WARMUP
        num_iters = NUM_ITERS_FAST if fast_mode else NUM_ITERS
        
        task = []
        gtimes = {}
        for solidx in self.ck_solutions:
            info = (
                (
                    self.input.shape,
                    self.w1_qweight.shape,
                    None if self.w1_scale is None else self.w1_scale.shape,
                    None if self.w1_zp is None else self.w1_zp.shape,
                    self.w2_qweight.shape,
                    None if self.w2_scale is None else self.w2_scale.shape,
                    None if self.w2_zp is None else self.w2_zp.shape,
                    self.topk_weights.shape,
                    self.topk_ids.shape,
                    [self.block_shape_n, self.block_shape_k],
                ),
                solidx,
            )

            task.append(
                (
                    info,
                    #generated data function set none if the input data has been generated
                    None,
                    #gen args
                    None,
                    run_fused_experts_ck_impl,
                    (
                        self.input,
                        self.w1_qweight,
                        self.w2_qweight,
                        self.topk_weights,
                        self.topk_ids,
                        self.input.dtype,
                        False,  # inplace
                        "silu",  # activation
                        self.use_f8_w8a8_block,
                        self.use_int8_w8a8_block,
                        self.use_int8_w8a16,
                        self.use_int4_w4a16,
                        self.use_int4_w4a8_block,
                        self.per_channel_quant,
                        self.expert,
                        self.block_m,
                        None,
                        self.w1_scale,
                        self.w2_scale,
                        self.w1_zp,
                        self.w2_zp,
                        self.a1_scale,
                        self.a2_scale,
                        [self.block_shape_n, self.block_shape_k],
                        False,
                        solidx,
                    ),
                    {
                        "num_warmup": num_warmup,
                        "num_iters": num_iters,
                    },
                    None, #ref func if is none, use precomputed golden instead
                    (), # ref args
                    {}, # ref kwargs
                    self.ref if fast_mode == 0 else None,
                    self.rtol,
                    self.atol,
                )
            )

        in_data = [
            (
                len(self.ck_solutions),
                #input data if all solutions use same input
                #if you want to use diffenent data,please pass this in task args posistion.
                (
                    self.input, self.w1_qweight, self.w2_qweight, self.topk_weights, self.topk_ids,
                ),
            )
        ]
            
        ret = mp_tuner(task, in_data, self.mp, fast_mode, False)
        for info, us, err_ratio in ret:
            if fast_mode == 0:
                if err_ratio > self.check_err_ratio:
                    continue
            solidx = info[-1]
            gtimes[solidx] = us / 1000.0
            print(f"solution id: {solidx}, time: {gtimes[solidx]:.3f} ms, err_ratio: {err_ratio:.6f}")
        self.ck_gtimedf = pd.DataFrame.from_dict(
            gtimes, orient="index", columns=["gtimems"]
        ).sort_values(by="gtimems")
        self.ck_gtimedf.to_csv("/tmp/ck_gtimedf.csv")
        print(">>> CK top solutions, Fast Mode", fast_mode, flush=True)
    
    
    def warmup(self, warmi=500):
        for i in range(warmi):
            self.blob = self.blob + 0.00001
            
    def find_ck_solutions(self, fast_mode=0):
        self.ck_solutions = ck_moe_get_solutions(self.input, self.w1_qweight, self.w2_qweight, self.topk_weights, self.topk_ids, 
                                                use_int8_w8a16 = self.use_int8_w8a16,
                                                use_int4_w4a16 = self.use_int4_w4a16,
                                                use_int8_w8a8_block = self.use_int8_w8a8_block,
                                                use_int4_w4a8_block = self.use_int4_w4a8_block,
                                                w1_zp = self.w1_zp,
                                                w2_zp = self.w2_zp,
                                                w1_scale = self.w1_scale,
                                                w2_scale = self.w2_scale,
                                                block_shape_n = self.block_shape_n,
                                                block_shape_k = self.block_shape_k,
                                                block_m = self.block_m)
        print(f"Found CK {len(self.ck_solutions)} solutions: {self.ck_solutions}")
        self.warmup()   # Is this necessary?
        self.ck_time_all_sols(fast_mode)
  

    def triton_time_all_sols(self, fast_mode=0):
        num_warmup = NUM_WARMUP_FAST if fast_mode else NUM_WARMUP
        num_iters = NUM_ITERS_FAST if fast_mode else NUM_ITERS
        pass

    def find_triton_solutions(self, fast_mode=0):
        # 1. get all solutions in triton


        # 2. time all solutions and pick the best one
        # self.warmup()
        self.triton_time_all_sols(fast_mode)


    def asm_time_all_sols(self, fast_mode=0):
        num_warmup = NUM_WARMUP_FAST if fast_mode else NUM_WARMUP
        num_iters = NUM_ITERS_FAST if fast_mode else NUM_ITERS
        task = []
        gtimes = {}
        solutions = 0
        for solidx in self.asm_solutions:
            info = (
                (
                    self.input.shape,
                    self.w1_qweight.shape,
                    None if self.w1_scale is None else self.w1_scale.shape,
                    None if self.w1_zp is None else self.w1_zp.shape,
                    self.w2_qweight.shape,
                    None if self.w2_scale is None else self.w2_scale.shape,
                    None if self.w2_zp is None else self.w2_zp.shape,
                    self.topk_weights.shape,
                    self.topk_ids.shape,
                    [self.block_shape_n, self.block_shape_k],
                ),
                solidx,
            )
            task.append(
                (
                    info,
                    #generated data function set none if the input data has been generated 
                    None,
                    #gen args
                    None,
                    run_fused_experts_asm_impl, # need to be wrapped by run func
                    (
                        self.input, 
                        self.w1_qweight, 
                        self.w2_qweight, 
                        self.topk_weights, 
                        self.topk_ids,
                        self.indtype,
                        False,  # inplace
                        "silu",  # activation
                        (self.use_f8_w8a8_block or self.use_f8_w8a8_channel),
                        (self.use_int8_w8a8_block or self.use_int8_w8a8_channel),
                        self.use_int4_w4a8_block,
                        self.use_int8_w8a16,  # use_int8_w8a16
                        self.use_int4_w4a16,
                        self.per_channel_quant,
                        self.expert,  # This should be an integer, not a tensor
                        None,  # expert_map
                        self.w1_scale,
                        self.w2_scale,
                        self.w1_zp,
                        self.w2_zp,
                        None,  # a1_scale
                        None,  # a2_scale
                        [self.block_shape_n, self.block_shape_k],
                        False,
                        0,
                        0, #shuffle
                        solidx,
                    ),
                    {
                        "num_warmup": num_warmup,
                        "num_iters": num_iters,
                    },
                    None, #ref func if is none, use precomputed golden instead
                    (), # ref args 
                    {}, # ref kwargs
                    self.ref if fast_mode == 0 else None, # ref golden
                    self.rtol,
                    self.atol,
                )
            )
            solutions = solutions +1
        in_data = [
            (
                solutions,
                #input data if all solutions use same input
                #if you want to use diffenent data,please pass this in task args posistion.
                (
                    self.input, self.w1_qweight, self.w2_qweight, self.topk_weights, self.topk_ids,                      
                ),
            )
        ]
        ret = mp_tuner(task,in_data,self.mp, fast_mode, False)
        for info, us, err_ratio in ret:
            if fast_mode == 0:
                if err_ratio > self.check_err_ratio:
                    continue
            solidx = info[-1]
            gtimes[solidx] = us / 1000.0
            print(f"solution id: {solidx}, time: {gtimes[solidx]:.3f} ms, err_ratio: {err_ratio:.6f}")
        self.asm_gtimedf = pd.DataFrame.from_dict(
            gtimes, orient="index", columns=["gtimems"]
        ).sort_values(by="gtimems")
        self.asm_gtimedf.to_csv("/tmp/asm_gtimedf.csv")
        print(">>> ASM top solutions, Fast Mode", fast_mode, flush=True)


    def find_asm_solutions(self, fast_mode=0):
        # 1. get all solutions in asm

        self.asm_solutions = asm_moe_get_solutions(self.input, self.w1_qweight, self.w2_qweight, self.topk_weights, self.topk_ids,
                                                use_int8_w8a16 = self.use_int8_w8a16,
                                                use_int4_w4a16 = self.use_int4_w4a16,
                                                use_int8_w8a8 = self.use_int8_w8a8_block or self.use_int8_w8a8_channel,
                                                use_int4_w4a8 = self.use_int4_w4a8_block,
                                                use_fp8_w8a8 = self.use_f8_w8a8_block or self.use_f8_w8a8_channel,
                                                per_channel_quant = self.per_channel_quant,
                                                w1_zp = self.w1_zp,
                                                w2_zp = self.w2_zp,
                                                w1_scale = self.w1_scale,
                                                w2_scale = self.w2_scale,
                                                block_shape_n = self.block_shape_n,
                                                block_shape_k = self.block_shape_k,
                                                block_m = self.block_m)
        print(f"Found ASM {len(self.asm_solutions)} solutions: {self.asm_solutions}")
        for combined in self.asm_solutions:
            parts = combined.split("+")
            if len(parts) == 2:
                key1 = int(parts[0])
                key2 = int(parts[1])
            else:
                raise ValueError("Invalid solution_id")
            print(f"{combined}\t\t{key1}\t\t{key2}")
        # 2. time all solutions and pick the best one
        # self.warmup()
        self.asm_time_all_sols(fast_mode)


    def find_fastest_solution(self):
        # 目前CK && asm，后续需比较triton
        if len(self.ck_gtimedf) > 0 and len(self.asm_gtimedf) > 0:
            if (self.ck_gtimedf.iloc[0,0]< self.asm_gtimedf.iloc[0,0]):
                self.sol_type = SolutionType.CK
                self.sol_id = int(self.ck_gtimedf.index[0])
                self.time_us = float(self.ck_gtimedf.iloc[0,0] * 1000)
            else:
                self.sol_type = SolutionType.ASM
                self.sol_id = self.asm_gtimedf.index[0]
                self.time_us = float(self.asm_gtimedf.iloc[0,0] * 1000)
        elif len(self.ck_gtimedf) > 0:
            self.sol_type = SolutionType.CK
            self.sol_id = int(self.ck_gtimedf.index[0])
            self.time_us = float(self.ck_gtimedf.iloc[0,0] * 1000)
        elif len(self.asm_gtimedf) > 0:
            self.sol_type = SolutionType.ASM
            self.sol_id = self.asm_gtimedf.index[0]
            self.time_us = float(self.asm_gtimedf.iloc[0,0] * 1000)

        if self.sol_type == SolutionType.ASM:
            parts = self.sol_id.split("+")
            key1 = int(parts[0])
            key2 = int(parts[1])
            print(f"Fastest solution: {self.sol_type}, id: {self.sol_id}={key1}+{key2}, time: {self.time_us:.3f} us")
        else:
            print(f"Fastest solution: {self.sol_type}, id: {self.sol_id}, time: {self.time_us:.3f} us")
