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

e2e_time_param_dict_list = [{"token_num": 1, "expert_num": 32, "topk": 5, "input_dtype": torch.int32},
                            {"token_num": 16, "expert_num": 32, "topk": 5, "input_dtype": torch.int32},
                            {"token_num": 32, "expert_num": 32, "topk": 5, "input_dtype": torch.int32},
                            {"token_num": 64, "expert_num": 32, "topk": 5, "input_dtype": torch.int32},
                            {"token_num": 512, "expert_num": 32, "topk": 5, "input_dtype": torch.int32},
                            {"token_num": 1024, "expert_num": 32, "topk": 5, "input_dtype": torch.int32},
                            {"token_num": 4096, "expert_num": 32, "topk": 5, "input_dtype": torch.int32},
                            {"token_num": 8192, "expert_num": 32, "topk": 5, "input_dtype": torch.int32},
                            {"token_num": 32767, "expert_num": 32, "topk": 5, "input_dtype": torch.int32},
                            {"token_num": 1, "expert_num": 8, "topk": 2, "input_dtype": torch.int32},
                            {"token_num": 16, "expert_num": 8, "topk": 2, "input_dtype": torch.int32},
                            {"token_num": 32, "expert_num": 8, "topk": 2, "input_dtype": torch.int32},
                            {"token_num": 64, "expert_num": 8, "topk": 2, "input_dtype": torch.int32},
                            {"token_num": 512, "expert_num": 8, "topk": 2, "input_dtype": torch.int32},
                            {"token_num": 1024, "expert_num": 8, "topk": 2, "input_dtype": torch.int32},
                            {"token_num": 4096, "expert_num": 8, "topk": 2, "input_dtype": torch.int32},
                            {"token_num": 8192, "expert_num": 8, "topk": 2, "input_dtype": torch.int32},
                            {"token_num": 32767, "expert_num": 8, "topk": 2, "input_dtype": torch.int32}]
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", "expert_num", "topk", "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"]
    expert_num = params_dict["expert_num"]
    topk = params_dict["topk"]
    dtype = params_dict["input_dtype"]
    expert_id = torch.randint(low=0, high=expert_num, size=(token_num, topk)).to(torch.int32).to('mlu')
    gather_idx = torch.empty((token_num * topk), dtype=dtype, device='mlu')
    combine_idx = torch.empty((token_num * topk), dtype=dtype, device='mlu')
    token_count = torch.empty((expert_num), dtype=dtype, device='mlu')
    cusum_token_count = torch.empty((expert_num + 1), dtype=dtype, device='mlu')
    hardware_time, e2e_time = benchmark_forward(tmo.moe_gen_idx,
                                                expert_id,
                                                expert_num,
                                                repeats=args.repeat_times)
    io_bytes = expert_id.element_size() * expert_id.nelement() + \
               gather_idx.element_size() * gather_idx.nelement() + \
               combine_idx.element_size() * combine_idx.nelement() + \
               token_count.element_size() * token_count.nelement() + \
               cusum_token_count.element_size() * cusum_token_count.nelement()
    io_eff = io_bytes / hardware_time / bd
    content = [f"{token_num}", f"{expert_num}", f"{topk}", 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()