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
from itertools import product

e2e_time_param_dict_list = [{"batch": 1,  "seq_len": 1024, "head_num": 25, "head_size": 64, "has_residual": True,
                             "has_bias": True, "has_quant": True, "eps": 1e-6, "dynamic_quant": [False, True], "input_dtype": [torch.float16, torch.bfloat16]},
                            {"batch": 1,  "seq_len": 1024, "head_num": 16, "head_size": 128, "has_residual": True,
                             "has_bias": True, "has_quant": True, "eps": 1e-6, "dynamic_quant": [False, True], "input_dtype": [torch.float16, torch.bfloat16]},
                            {"batch": 1,  "seq_len": 1024, "head_num": 32, "head_size": 128, "has_residual": True,
                             "has_bias": True, "has_quant": True, "eps": 1e-6, "dynamic_quant": [False, True], "input_dtype": [torch.float16, torch.bfloat16]},
                            {"batch": 1,  "seq_len": 1024, "head_num": 40, "head_size": 128, "has_residual": True,
                             "has_bias": True, "has_quant": True, "eps": 1e-6, "dynamic_quant": [False, True], "input_dtype": [torch.float16, torch.bfloat16]},
                            {"batch": 1,  "seq_len": 1024, "head_num": 64, "head_size": 96, "has_residual": True,
                             "has_bias": True, "has_quant": True, "eps": 1e-6, "dynamic_quant": [False, True], "input_dtype": [torch.float16, torch.bfloat16]},
                            {"batch": 1,  "seq_len": 1024, "head_num": 52, "head_size": 128, "has_residual": True,
                             "has_bias": True, "has_quant": True, "eps": 1e-6, "dynamic_quant": [False, True], "input_dtype": [torch.float16, torch.bfloat16]},
                            {"batch": 1,  "seq_len": 1024, "head_num": 64, "head_size": 128, "has_residual": True,
                             "has_bias": True, "has_quant": True, "eps": 1e-6, "dynamic_quant": [False, True], "input_dtype": [torch.float16, torch.bfloat16]},
                            {"batch": 1,  "seq_len": 1024, "head_num": 96, "head_size": 128, "has_residual": True,
                             "has_bias": True, "has_quant": True, "eps": 1e-6, "dynamic_quant": [False, True], "input_dtype": [torch.float16, torch.bfloat16]},
                            {"batch": 1,  "seq_len": 1024, "head_num": 112, "head_size": 128, "has_residual": True,
                             "has_bias": True, "has_quant": True, "eps": 1e-6, "dynamic_quant": [False, True], "input_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 = ["input_shape", "has_residual", "has_bias", "has_quant", "input_dtype", "hardware_time(us)", "e2e_latency(us)"]
  contents = []
  for params_dict in e2e_time_param_dict_list:
    batch = params_dict["batch"]
    seq_len = params_dict["seq_len"]
    head_num = params_dict["head_num"]
    head_size = params_dict["head_size"]
    has_residual = params_dict["has_residual"]
    has_quant = params_dict["has_quant"]
    has_bias = params_dict["has_bias"]
    eps = params_dict["eps"]
    dynamic_quant_list = params_dict["dynamic_quant"]
    input_dtype_list = params_dict["input_dtype"]
    dynamic_quant_list = params_dict["dynamic_quant"]
    input_dtype_list = params_dict["input_dtype"]
    iters = product(dynamic_quant_list, input_dtype_list)
    for dynamic_quant, dtype in iters:
      if dtype == torch.bfloat16 and not torch_mlu.mlu.is_bf16_supported():
        continue
      x = torch.randn(batch, seq_len, head_num, head_size).to(dtype).to(device)
      beta = torch.randn(head_size).to(dtype).to(device)
      gamma = torch.randn(head_size).to(dtype).to(device)
      residual, bias, quant_scale = None, None, None
      if has_residual:
        residual = torch.randn(batch, seq_len, head_num, head_size).to(dtype).to(device)
      if has_bias:
        bias = torch.randn(head_size).to(dtype).to(device)
      if has_quant or dynamic_quant:
        quant_scale = torch.randn(head_size).to(device)
      hardware_time, e2e_time = benchmark_forward(tmo.fused_rms_norm,
                                                  x,
                                                  residual,
                                                  gamma,
                                                  beta,
                                                  bias,
                                                  eps,
                                                  False,
                                                  quant_scale,
                                                  None,
                                                  dynamic_quant,
                                                  repeats=args.repeat_times)
      content = [f"{batch, seq_len, head_num, head_size}", f"{has_residual}", f"{has_bias}", f"{has_quant}", 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()
