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
import random

e2e_time_param_dict_list = [{"batch": 16, "head_num": 8, "head_size": 128, "block_seq_len": 1, "max_seq_len": 1,
                             "dtype": [torch.float16], "pack": False},
                            {"batch": 128, "head_num": 8, "head_size": 128, "block_seq_len": 1, "max_seq_len": 1,
                             "dtype": [torch.float16], "pack": False},
                            {"batch": 512, "head_num": 8, "head_size": 128, "block_seq_len": 1, "max_seq_len": 1,
                             "dtype": [torch.float16], "pack": False},
                            {"batch": 1024, "head_num": 8, "head_size": 128, "block_seq_len": 1, "max_seq_len": 1,
                             "dtype": [torch.float16], "pack": False},
                            {"batch": 1, "head_num": 8, "head_size": 128, "block_seq_len": 1024, "max_seq_len": 1024,
                             "dtype": [torch.float16], "pack": True},
                            {"batch": 1, "head_num": 8, "head_size": 128, "block_seq_len": 2048, "max_seq_len": 2048,
                             "dtype": [torch.float16], "pack": True},
                            {"batch": 1, "head_num": 8, "head_size": 128, "block_seq_len": 8192, "max_seq_len": 8192,
                             "dtype": [torch.float16], "pack": True},
                            {"batch": 1, "head_num": 8, "head_size": 128, "block_seq_len": 32768, "max_seq_len": 32768,
                             "dtype": [torch.float16], "pack": True},]

def gen_tenosr(batch, head_num, head_size, block_seq_len, max_seq_len, dtype, pack):
    if not pack:
        out           = torch.randn(batch, max_seq_len, head_num, head_size, device="mlu", dtype=dtype)
        lse           = torch.randn(batch, head_num, max_seq_len, device="mlu", dtype=torch.float32)
        block_out     = torch.randn(batch, block_seq_len, head_num, head_size, device="mlu", dtype=dtype)
        block_lse     = torch.randn(batch, head_num, block_seq_len, device="mlu", dtype=torch.float32)
        seq_offset    = None
        cu_seqs       = None
        block_cu_seqs = None
    else:
        seq_lens         = torch.randint(low=max_seq_len, high=(max_seq_len + 1), size=(batch, ), dtype=torch.int32)
        block_seq_lens   = torch.randint(low=block_seq_len, high=(block_seq_len + 1), size=(batch, ), dtype=torch.int32)
        block_seq_lens   = torch.minimum(seq_lens, block_seq_lens)
        seq_offset       = torch.zeros_like(seq_lens)
        for i in range(batch):
            seq_offset[i] = torch.randint(low=0, high=seq_lens[i]-block_seq_lens[i]+1, size=(1,), dtype=torch.int32)
        seq_offset       = seq_offset.mlu()
        cu_seqs          = torch.cat((torch.tensor([0]),torch.cumsum(seq_lens, dim=0))).to(torch.int32).mlu()
        block_cu_seqs    = torch.cat((torch.tensor([0]),torch.cumsum(block_seq_lens, dim=0))).to(torch.int32).mlu()
        total_seqs       = torch.sum(seq_lens)
        block_total_seqs = torch.sum(block_seq_lens)

        out           = torch.randn(total_seqs, head_num, head_size, device="mlu", dtype=dtype)
        lse           = torch.randn(batch, head_num, max_seq_len, device="mlu", dtype=torch.float32)
        block_out     = torch.randn(block_total_seqs, head_num, head_size, device="mlu", dtype=dtype)
        block_lse     = torch.randn(batch, head_num, block_seq_len, device="mlu", dtype=torch.float32)
    return (out, lse, block_out, block_lse, seq_offset, cu_seqs, block_cu_seqs)


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", "head_num", "head_size", "block_seq_len", "max_seq_len", "dtype", "pack", "hardware_time(us)", "e2e_latency(us)"]
  contents = []
  for params_dict in e2e_time_param_dict_list:
    batch = params_dict["batch"]
    head_num = params_dict["head_num"]
    head_size = params_dict["head_size"]
    block_seq_len = params_dict["block_seq_len"]
    max_seq_len = params_dict["max_seq_len"]
    dtype_list = params_dict["dtype"]
    pack = params_dict["pack"]
    for dtype in dtype_list:
      out, lse, block_out, block_lse, seq_offset, cu_seqs, block_cu_seqs =    \
        gen_tenosr(batch, head_num, head_size, block_seq_len, max_seq_len, dtype, pack)

      hardware_time, e2e_time = benchmark_forward(tmo.update_out_and_lse,
                                                  out,
                                                  lse,
                                                  block_out,
                                                  block_lse,
                                                  seq_offset,
                                                  cu_seqs,
                                                  block_cu_seqs,
                                                  repeats=args.repeat_times)
      content = [f"{batch}", f"{head_num}", f"{head_size}", f"{block_seq_len}", f"{max_seq_len}", f"{dtype}", f"{pack}", 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()
