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

e2e_time_param_dict_list = [{"batch": 1, "seq_len": 2048, "hidden_size": 8192, "expert_num": 32, "input_dtype": torch.bfloat16},
                            {"batch": 1, "seq_len": 4096, "hidden_size": 8192, "expert_num": 32, "input_dtype": torch.bfloat16},
                            {"batch": 1, "seq_len": 32768, "hidden_size": 8192, "expert_num": 32, "input_dtype": torch.bfloat16},
                            {"batch": 16, "seq_len": 1, "hidden_size": 8192, "expert_num": 32, "input_dtype": torch.bfloat16},
                            {"batch": 128, "seq_len": 1, "hidden_size": 8192, "expert_num": 32, "input_dtype": torch.bfloat16},
                            {"batch": 512, "seq_len": 1, "hidden_size": 8192, "expert_num": 32, "input_dtype": torch.bfloat16}]

def main():
  if 'MLU3' in torch.mlu.get_device_name():
    exit()
  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 = ["batch", "seq_len", "hidden_size", "expert_num", "input_dtype", "hardware_time(us)",
            "e2e_latency(us)", "IO efficiency"]

  contents = []
  bandwidth = get_band_width()
  for param_dict in e2e_time_param_dict_list:
    batch = param_dict["batch"]
    seq_len = param_dict["seq_len"]
    hidden_size = param_dict["hidden_size"]
    expert_num = param_dict["expert_num"]
    input_dtype = param_dict["input_dtype"]
    if input_dtype == torch.bfloat16 and not torch_mlu.mlu.is_bf16_supported():
        input_dtype = torch.half
    input = torch.randn(batch, seq_len, hidden_size, dtype=input_dtype, device="mlu")
    weight = torch.randn(expert_num, hidden_size, dtype=torch.float32, device="mlu")
    hardware_time, e2e_time = benchmark_forward(tmo.moe_cast_gating,
                                                input,
                                                weight)
    io_bytes = batch * seq_len * hidden_size * input.element_size() + \
        expert_num * hidden_size * weight.element_size() + batch * seq_len * expert_num * weight.element_size()
    io_coeff = io_bytes / hardware_time / bandwidth
    content = [f"{batch}", f"{seq_len}", f"{hidden_size}", f"{expert_num}", f"{input_dtype}",
               f"{hardware_time}", f"{e2e_time}", f"{io_coeff}"]
    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()