import torch
import torch_mlu
import torch_mlu_ops as tmo
from common import *
import argparse
from tabulate import tabulate
import os
import numpy as np

e2e_time_param_dict_list = [{"token_num": 1, "hidden_size": 4096, "expert_num": 32, "topk": 5,
                             "start_expert_id": 0, "expert_size": 32, "input_dtype": [torch.int8, torch.float16, torch.bfloat16]},
                            {"token_num": 16, "hidden_size": 4096, "expert_num": 32, "topk": 5,
                             "start_expert_id": 0, "expert_size": 32, "input_dtype": [torch.int8, torch.float16, torch.bfloat16]},
                            {"token_num": 32, "hidden_size": 4096, "expert_num": 32, "topk": 5,
                             "start_expert_id": 0, "expert_size": 32, "input_dtype": [torch.int8, torch.float16, torch.bfloat16]},
                            {"token_num": 64, "hidden_size": 4096, "expert_num": 32, "topk": 5,
                             "start_expert_id": 0, "expert_size": 32, "input_dtype": [torch.int8, torch.float16, torch.bfloat16]},
                            {"token_num": 128, "hidden_size": 4096, "expert_num": 32, "topk": 5,
                             "start_expert_id": 0, "expert_size": 32, "input_dtype": [torch.int8, torch.float16, torch.bfloat16]},
                            {"token_num": 512, "hidden_size": 4096, "expert_num": 32, "topk": 5,
                             "start_expert_id": 0, "expert_size": 32, "input_dtype": [torch.int8, torch.float16, torch.bfloat16]},
                            {"token_num": 1024, "hidden_size": 4096, "expert_num": 32, "topk": 5,
                             "start_expert_id": 0, "expert_size": 32, "input_dtype": [torch.int8, torch.float16, torch.bfloat16]},
                            {"token_num": 4096, "hidden_size": 4096, "expert_num": 32, "topk": 5,
                             "start_expert_id": 0, "expert_size": 32, "input_dtype": [torch.int8, torch.float16, torch.bfloat16]},
                            {"token_num": 8192, "hidden_size": 4096, "expert_num": 32, "topk": 5,
                             "start_expert_id": 0, "expert_size": 32, "input_dtype": [torch.int8, torch.float16, torch.bfloat16]},
                            {"token_num": 32768, "hidden_size": 4096, "expert_num": 32, "topk": 5,
                             "start_expert_id": 0, "expert_size": 32, "input_dtype": [torch.int8, torch.float16, torch.bfloat16]}]

def gen_tensor(token_num, hidden_size, expert_num, topk, start_expert_id, expert_size, dtype):
    input = torch.randn(token_num, hidden_size).to(dtype).to('mlu')
    gather_idx = torch.randint(low=0, high=token_num, size=(token_num * topk,)).to(torch.int32).to('mlu')
    cusum_token_count, _ = generate_token_count(expert_num, token_num * topk)
    cusum_token_count = cusum_token_count.to('mlu')
    use_all_experts = expert_num == expert_size
    if use_all_experts:
      cusum_token_count = None
      real_token_count = token_num * topk
    else:
      real_token_count = cusum_token_count[start_expert_id+expert_size] - cusum_token_count[start_expert_id]
    return input, gather_idx, cusum_token_count, real_token_count

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--repeat_times', type=int, default=10, help='repeat times for testing')
    parser.add_argument('--csv', action='store_true', help='write the report data to csv')
    parser.add_argument('-o', type=str, help='specify the output folder name under --csv mode')

    args = parser.parse_args()

    titles = ["token_num", "hidden_size", "expert_num", "topk", "start_expert_id", "expert_size", "input_dtype",
              "hardware_time(us)", "e2e_latency(us)", "IO efficiency"]
    contents = []
    bd = get_band_width()
    for params_dict in e2e_time_param_dict_list:
        token_num = params_dict["token_num"]
        hidden_size = params_dict["hidden_size"]
        expert_num = params_dict["expert_num"]
        topk = params_dict["topk"]
        start_expert_id = params_dict["start_expert_id"]
        expert_size = params_dict["expert_size"]
        input_dtype_list = params_dict["input_dtype"]
        for dtype in input_dtype_list:
            if dtype == torch.bfloat16 and not torch_mlu.mlu.is_bf16_supported():
              continue
            input, gather_idx, cusum_token_count, real_token_count = \
              gen_tensor(token_num, hidden_size, expert_num,topk, start_expert_id, expert_size, dtype)
            hardware_time, e2e_time = benchmark_forward(tmo.moe_expand_input,
                                                        input,
                                                        gather_idx,
                                                        cusum_token_count,
                                                        start_expert_id,
                                                        expert_size,
                                                        repeats=args.repeat_times)
            io_bytes = input.element_size() * input.nelement() + \
                      gather_idx.element_size() * gather_idx.nelement() + \
                      (cusum_token_count.element_size() * cusum_token_count.nelement() if cusum_token_count is not None else 0) + \
                      real_token_count * input.element_size()
            io_eff = io_bytes / hardware_time / bd
            content = [f"{token_num}", f"{hidden_size}", f"{expert_num}", f"{topk}", f"{start_expert_id}", f"{expert_size}",
                      f"{dtype}", f"{hardware_time}", f"{e2e_time}", f"{io_eff}"]
            contents.append(content)
    table = [titles] + contents
    print(tabulate(table, headers="firstrow", tablefmt="grid"))

    if args.csv:
        current_file_path = __file__
        _, file_name = os.path.split(current_file_path)
        save_to_csv(table, args.o, file_name)

if __name__ == "__main__":
  main()