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

params_dict = {
                "token_num": [n * 5 for n in [1, 72, 512, 1024, 4096, 32768]],
                "hidden_size": [1024, 8192],
                "input_dtype": [torch.float16]
              }

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()
  device = 'mlu'
  titles = ["token_num", "hidden_size", "input_dytpe", "hardware_time(us)", "e2e_latency(us)", "IO efficiency"]
  contents = []
  params_list = product(params_dict["token_num"], params_dict["hidden_size"], params_dict["input_dtype"])
  bd = get_band_width()
  for params in params_list:
    token_num, hidden_size = params[0], params[1]
    input_shape = (token_num, hidden_size)
    smooth_shape = (hidden_size)
    dtype = params[2]
    if dtype == torch.bfloat16 and not torch_mlu.mlu.is_bf16_supported():
        continue
    input = torch.randn(input_shape).to(device).to(dtype)
    smooth = torch.randn(smooth_shape).to(device).to(torch.float32)
    zero = None
    token_count = None
    hardware_time, e2e_time = benchmark_forward(tmo.per_token_smooth_quantize,
                                                input,
                                                smooth,
                                                zero,
                                                token_count,
                                                repeats=args.repeat_times)
    io_bytes = (input.element_size() + 1) * input.nelement() + \
                smooth.element_size() * smooth.nelement() + \
                token_num * 4
    io_eff = io_bytes / hardware_time / bd
    content = [f"{token_num}", f"{hidden_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()
