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 = [{"batch": 1, "seq_len": 1024, "input_size": 1600, "head_size": 80, "hidden_size": 1600, "input_dtype": [torch.float16, torch.bfloat16]},
                            {"batch": 1, "seq_len": 1024, "input_size": 2048, "head_size": 128, "hidden_size": 2048, "input_dtype": [torch.float16, torch.bfloat16]},
                            {"batch": 1, "seq_len": 1024, "input_size": 4096, "head_size": 128, "hidden_size": 4096, "input_dtype": [torch.float16, torch.bfloat16]},
                            {"batch": 1, "seq_len": 1024, "input_size": 6144, "head_size": 128, "hidden_size": 6144, "input_dtype": [torch.float16, torch.bfloat16]},
                            {"batch": 1, "seq_len": 1024, "input_size": 6656, "head_size": 128, "hidden_size": 6656, "input_dtype": [torch.float16, torch.bfloat16]},
                            {"batch": 1, "seq_len": 1024, "input_size": 8192, "head_size": 128, "hidden_size": 8192, "input_dtype": [torch.float16, torch.bfloat16]},
                            {"batch": 1, "seq_len": 1024, "input_size": 12288, "head_size": 128, "hidden_size": 12288, "input_dtype": [torch.float16, torch.bfloat16]},
                            {"batch": 1, "seq_len": 1024, "input_size": 14336, "head_size": 128, "hidden_size": 14336, "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 = ["batch", "seq_len", "input_size", "hidden_size", "head_size", "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"]
    input_size = params_dict["input_size"]
    hidden_size = params_dict["hidden_size"]
    head_size = params_dict["head_size"]
    input_dtype_list = params_dict["input_dtype"]
    for dtype in input_dtype_list:
      if dtype == torch.bfloat16 and not torch_mlu.mlu.is_bf16_supported():
        continue
      input = torch.randn(batch, seq_len, input_size).to(device).to(dtype)
      weight = torch.randn(hidden_size * 3, input_size).to(device).to(dtype)
      bias = torch.randn(hidden_size * 3).to(device).to(dtype)
      weights = torch.chunk(weight, 3)
      biases = torch.chunk(bias, 3)
      norm_weight = torch.randn(input_size).to(device).to(dtype)
      norm_bias = torch.randn(input_size).to(device).to(dtype)
      hardware_time, e2e_time = benchmark_forward(tmo.fused_norm_attention_project,
                                                  input,
                                                  weights[0],
                                                  biases[0],
                                                  weights[1],
                                                  biases[1],
                                                  weights[2],
                                                  biases[2],
                                                  norm_weight,
                                                  norm_bias,
                                                  1e-6,
                                                  'nthc',
                                                  head_size,
                                                  False,
                                                  repeats=args.repeat_times)
      content = [f"{batch}", f"{seq_len}", f"{input_size}", f"{hidden_size}", f"{head_size}", 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()
