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

e2e_time_param_dict_list = [
                            {"batch": 1,  "seq_len": 2048,  "hidden_size": 8192, "has_residual": True,
                             "has_bias": False, "has_quant": True, "dynamic_quant": True,
                             "input_dtype": torch.bfloat16},
                            {"batch": 1,  "seq_len": 4096,  "hidden_size": 8192, "has_residual": True,
                             "has_bias": False, "has_quant": True, "dynamic_quant": True,
                             "input_dtype": torch.bfloat16},
                            {"batch": 1,  "seq_len": 8192,  "hidden_size": 8192, "has_residual": True,
                             "has_bias": False, "has_quant": True, "dynamic_quant": True,
                             "input_dtype": torch.bfloat16},
                            {"batch": 1,  "seq_len": 32768,  "hidden_size": 8192, "has_residual": True,
                             "has_bias": False, "has_quant": True, "dynamic_quant": True,
                             "input_dtype": torch.bfloat16},
                            {"batch": 490,  "seq_len": 1,  "hidden_size": 8192, "has_residual": True,
                             "has_bias": False, "has_quant": True, "dynamic_quant": True,
                             "input_dtype": torch.bfloat16},
                            {"batch": 525,  "seq_len": 1,  "hidden_size": 8192, "has_residual": True,
                             "has_bias": False, "has_quant": True, "dynamic_quant": True,
                             "input_dtype": 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", "hidden_size", "has_residual", "has_bias", "has_quant",
            "dynamic_quant", "input_dtype", "hardware_time(us)", "e2e_latency(us)", "IO efficiency"]
  contents = []
  bd = get_band_width()
  for params_dict in e2e_time_param_dict_list:
    batch = params_dict["batch"]
    seq_len = params_dict["seq_len"]
    hidden_size = params_dict["hidden_size"]
    has_residual = params_dict["has_residual"]
    has_bias = params_dict["has_bias"]
    has_quant = params_dict["has_quant"]
    dynamic_quant = params_dict["dynamic_quant"]
    dtype = params_dict["input_dtype"]
    eps = 1e-6

    if dtype == torch.bfloat16 and not torch_mlu.mlu.is_bf16_supported():
      dtype = torch.half

    x = torch.randn(batch, seq_len, hidden_size, dtype=dtype, device=device)
    beta = torch.randn(hidden_size, dtype=dtype, device=device)
    gamma = torch.randn(hidden_size, dtype=dtype, device=device)
    residual, bias, quant_scale = None, None, None
    if has_residual:
      residual = torch.randn(batch, seq_len, hidden_size, dtype=dtype, device=device)
    if has_bias:
      bias = torch.randn(hidden_size, dtype=dtype, device=device)
    if has_quant or dynamic_quant:
      quant_scale = torch.randn(hidden_size, dtype=torch.float, device=device)

    store_output_before_norm = has_residual
    hardware_time, e2e_time = benchmark_forward(tmo.fused_layer_norm,
                                                x,
                                                residual,
                                                gamma,
                                                beta,
                                                bias,
                                                eps,
                                                store_output_before_norm,
                                                quant_scale,
                                                None,
                                                dynamic_quant,
                                                repeats=args.repeat_times)
    n = x.nelement()
    sizeoft = x.element_size()
    io_bytes = (sizeoft + 1) * n + \
               (1 + store_output_before_norm) * (sizeoft * n if has_residual else 0) + \
               sizeoft * hidden_size * 2 + \
               (hidden_size * 4 if has_quant else 0) + \
               (batch * seq_len * 4 if dynamic_quant else 0)
    io_eff = io_bytes / hardware_time / bd

    content = [f"{batch}", f"{seq_len}", f"{hidden_size}", f"{has_residual}", f"{has_bias}",
                f"{has_quant}", f"{dynamic_quant}", 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()
