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


e2e_time_param_dict_list = [
                            {"token_num": [1024, 2048, 3072, 4096], "head_num_kv": 1, "head_size": 128, "block_size": 16,
                             "input_dtype": [torch.float16, torch.bfloat16]},
                            {"token_num": [1024 * 32, 2048 * 32, 3072 * 32, 4096 * 32], "head_num_kv": 1, "head_size": 128, "block_size": 16,
                             "input_dtype": [torch.float16, torch.bfloat16]},
                            {"token_num": [1024 * 64, 2048 * 64, 3072 * 64, 4096 * 64], "head_num_kv": 1, "head_size": 128, "block_size": 16,
                             "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()
  titles = ["token_num", "head_num_kv", "head_size", "block_size", "input_dtype", "hardware_time(us)", "e2e_latency(us)"]
  contents = []
  if "MLU3" in torch.mlu.get_device_name():
    print("Op offline_quant_to_paged_cache does not support MLU300 devices.")
    return
  for params_dict in e2e_time_param_dict_list:
    token_num_list = params_dict["token_num"]
    # block_num = params_dict["block_num"]
    block_size = params_dict["block_size"]
    head_num_kv = params_dict["head_num_kv"]
    head_size = params_dict["head_size"]
    input_dtype_list = params_dict["input_dtype"]
    for token_num, dtype in product(token_num_list, input_dtype_list):
      if dtype == torch.bfloat16 and not torch_mlu.mlu.is_bf16_supported():
        continue
      block_num = (token_num + block_size - 1) // block_size
      key = torch.randn(token_num, head_num_kv, head_size, dtype=dtype, device="mlu")
      value = torch.randn(token_num, head_num_kv, head_size, dtype=dtype, device="mlu")
      key_cache = torch.randint(-128, 127, (block_num, head_num_kv, block_size, head_size), dtype=torch.int8).to("mlu")
      value_cache = torch.randint(-128, 127, (block_num, head_num_kv, block_size, head_size), dtype=torch.int8).to("mlu")
      num_slots = block_num * block_size
      slot_mapping = random.sample(range(num_slots), token_num)
      slot_mapping = torch.tensor(slot_mapping, dtype=torch.int32, device="mlu")
      key_cache_scale = torch.randn(head_num_kv, head_size, dtype=torch.float, device="mlu")
      value_cache_scale = torch.randn(head_num_kv, head_size, dtype=torch.float, device="mlu")
      hardware_time, e2e_time = benchmark_forward(tmo.offline_quant_to_paged_cache,
                                                  key, value,
                                                  key_cache_scale, value_cache_scale,
                                                  slot_mapping,
                                                  key_cache, value_cache,
                                                  repeats=args.repeat_times)


      content = [f"{token_num}", f"{head_num_kv}", f"{head_size}", f"{block_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()
