import torch
import torch_mlu
import torch_mlu_ops as tmo
from common import benchmark_forward, save_to_csv
import argparse
from tabulate import tabulate
import os

e2e_time_param_dict_list = [{"m": 1024, "k": 4096, "n": 14336, "has_c": True, "has_bias": True,
                             "act_mode": "none", "output_dtype": [torch.float16, torch.bfloat16]},
                            {"m": 1024, "k": 5120, "n": 13824, "has_c": False, "has_bias": True,
                             "act_mode": "silu", "output_dtype": [torch.float16, torch.bfloat16]}]

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 = ["m", "k", "n", "has_c", "has_bias", "act_mode", "output_dtype", "hardware_time(us)", "e2e_latency(us)"]
  contents = []
  for params_dict in e2e_time_param_dict_list:
    m = params_dict["m"]
    k = params_dict["k"]
    n = params_dict["n"]
    has_c = params_dict["has_c"]
    has_bias = params_dict["has_bias"]
    act_mode = params_dict["act_mode"]
    output_dtype_list = params_dict["output_dtype"]
    for dtype in output_dtype_list:
      if dtype == torch.bfloat16 and not torch_mlu.mlu.is_bf16_supported():
        continue
      a = torch.randn(m, k).to(device).to(torch.int8)
      b = torch.randn(n, k).to(device).to(torch.int8)
      a_scale = torch.randn(m).to(device)
      b_scale = torch.randn(n).to(device)
      c = None
      if has_c:
        c = torch.randn(m, n).to(device).to(dtype)
      bias = None
      if has_bias:
        bias = torch.randn(n).to(device).to(dtype)
      hardware_time, e2e_time = benchmark_forward(tmo.smooth_quant_matmul,
                                                  a,
                                                  a_scale,
                                                  b,
                                                  b_scale,
                                                  dtype,
                                                  bias,
                                                  c,
                                                  act_mode,
                                                  1.0,
                                                  1.0,
                                                  repeats=args.repeat_times)
      content = [f"{m}", f"{k}", f"{n}", f"{has_c}", f"{has_bias}", f"{act_mode}", f"{dtype}", f"{hardware_time}", f"{e2e_time}"]
      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()