import torch
import itertools
from aiter.test_common import perftest
from aiter.fused_moe_ck import MoeQuantType,get_moe_ck_solution
from aiter.fused_moe import fused_topk


BLOCK_SIZE_M = 16

@perftest(testGraph=True)
def ck_moe_test(
    ck_moe_sol,
    hidden_states,
    w1,
    w2,
    topk_weight,
    topk_ids,
    # following for int8 quant
    use_int8_w8a16 = False,
    use_int4_w4a16 = False,
    use_int8_w8a8_block = False,
    use_int4_w4a8_block = False,
    w1_zp = None,
    w2_zp = None,
    w1_scale=None,  # [expert, inter_dim, 1]
    w2_scale=None,  # [expert, model_dim, 1]
    fc1_smooth_scale=None,  # [expert, 1, model_dim]
    fc2_smooth_scale=None,  # [expert, 1, inter_dim]
    block_shape_n = 1,
    block_shape_k = 1
):
    return ck_moe_sol(
            hidden_states,
            w1,
            w2,
            topk_weight,
            topk_ids,
            use_int8_w8a16,
            use_int4_w4a16,
            use_int8_w8a8_block,
            use_int4_w4a8_block,
            w1_zp,
            w2_zp,
            w1_scale,
            w2_scale,
            fc1_smooth_scale,
            fc2_smooth_scale,
            block_shape_n,
            block_shape_k,
            block_m = BLOCK_SIZE_M
        ).to(hidden_states.dtype)



def moe_sol_test():
    M = [8, 16, 128]
    N = [256]
    K = [7168]
    E = [256]
    TOPK = [8]
    EP_SIZE = [1]
    DTYPE = [torch.float16]
    
    for m,n,k,e,topk,ep_size,dtype, in  itertools.product(
        M, N, K, E, TOPK, EP_SIZE, DTYPE):
        
        input = torch.randn((m, k), device="cuda", dtype=dtype) / 10
        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)
        
        ## without token topk score calc
        topk_weights, topk_ids = fused_topk(input, score, topk, True)
        
        
        # get ck solution
        ck_moe_sol = get_moe_ck_solution(dtype, m, n, k, e, topk, MoeQuantType.NO_QUANT)
        
        # performace test, in actual use case, we just call ck_moe_sol
        ck_out, ck_avg = ck_moe_test(ck_moe_sol, input, w1, w2, topk_weights, topk_ids)

        msg = f"[perf] token={m}, model_dim={k}, inter_dim={n}, {e=}, {topk=}, dtype: {dtype}, ck_avg:{ck_avg:<8.2f} us. == {ck_avg/1000:<6.4f} ms .... "
        print(msg)


if __name__ == "__main__":
    moe_sol_test()
