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

e2e_time_param_dict_list = [
                            {"batch": 16, "max_seq_len": 4096, "head_num_q": 8, "head_num_kv": 1, "head_size": 128,
                             "block_size": 16, "alibi_bias": False, "kv_cache_dtype": torch.int8, "use_paged_attn": False,
                             "is_pertoken": False, "input_dtype": [torch.bfloat16]},
                            {"batch": 128, "max_seq_len": 4096, "head_num_q": 8, "head_num_kv": 1, "head_size": 128,
                             "block_size": 16, "alibi_bias": False, "kv_cache_dtype": torch.int8, "use_paged_attn": False,
                             "is_pertoken": False, "input_dtype": [torch.bfloat16]},
                            {"batch": 512, "max_seq_len": 4096, "head_num_q": 8, "head_num_kv": 1, "head_size": 128,
                             "block_size": 16, "alibi_bias": False, "kv_cache_dtype": torch.int8, "use_paged_attn": False,
                             "is_pertoken": False, "input_dtype": [torch.bfloat16]},

                            {"batch": 16, "max_seq_len": 32768, "head_num_q": 8, "head_num_kv": 1, "head_size": 128,
                             "block_size": 16, "alibi_bias": False, "kv_cache_dtype": torch.int8, "use_paged_attn": False,
                             "is_pertoken": False, "input_dtype": [torch.bfloat16]},
                            {"batch": 128, "max_seq_len": 32768, "head_num_q": 8, "head_num_kv": 1, "head_size": 128,
                             "block_size": 16, "alibi_bias": False, "kv_cache_dtype": torch.int8, "use_paged_attn": False,
                             "is_pertoken": False, "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", "max_seq_len", "head_num_q", "head_num_kv", "head_size", "block_size", "alibi_bias", "kv_cache_dtype", "use_paged_attn", "input_dtype", "hardware_time(us)", "e2e_latency(us)"]
  contents = []
  for params_dict in e2e_time_param_dict_list:
    batch = params_dict["batch"]
    max_seq_len = params_dict["max_seq_len"]
    head_num_q = params_dict["head_num_q"]
    head_num_kv = params_dict["head_num_kv"]
    head_size = params_dict["head_size"]
    block_size = params_dict["block_size"]
    alibi_bias = params_dict["alibi_bias"]
    kv_cache_dtype = params_dict["kv_cache_dtype"]
    use_paged_attn = params_dict["use_paged_attn"]
    input_dtype_list = params_dict["input_dtype"]
    is_pertoken = params_dict["is_pertoken"]
    for dtype in input_dtype_list:
        if dtype == torch.bfloat16 and not torch_mlu.mlu.is_bf16_supported():
            continue
        input_qkv = torch.randn((batch, 1, head_num_q + 2 * head_num_kv, head_size)).to(device).to(dtype)
        input_q = input_qkv[..., 0 : head_num_q, :]

        context_lens = torch.randint(max_seq_len, max_seq_len + 1, (batch, ), dtype=torch.int32).to(device)
        max_context_len = int(max(context_lens))
        if use_paged_attn:
            mlu_name = torch.mlu.get_device_name()
            if "MLU3" in mlu_name:
                print("pagedattn is not implement on mlu370, skip it")
                continue
            block_size = 16
        else:
            block_size = max_seq_len + 512
        num_blocks = batch * ((max_seq_len + block_size - 1) // block_size)
        max_num_blocks_per_seq = (max_context_len + block_size - 1) // block_size
        cache_shape = (num_blocks, head_num_kv, block_size, head_size)
        scale_shape = (num_blocks, head_num_kv, block_size) if is_pertoken else (head_num_kv, head_size)
        block_tables = random.sample(range(0, num_blocks), batch * max_num_blocks_per_seq)
        block_tables = torch.tensor(block_tables, dtype=torch.int32).mlu().view(batch, max_num_blocks_per_seq)
        if kv_cache_dtype is not torch.int8:
            key_cache = torch.randn(size=cache_shape, dtype=torch.float16).to(device)
            value_cache = torch.randn(size=cache_shape, dtype=torch.float16).to(device)
            key_cache_scale = None
            value_cache_scale = None
        else:
            key_cache = torch.zeros(cache_shape).uniform_(-128, 128).to(kv_cache_dtype).to(device)
            value_cache = torch.zeros(cache_shape).uniform_(-128, 128).to(kv_cache_dtype).to(device)
            key_cache_scale = torch.randn(size=scale_shape, dtype=torch.float32).to(device)
            value_cache_scale = torch.randn(size=scale_shape, dtype=torch.float32).to(device)
        alibi_slopes = None
        if alibi_bias:
            alibi_slopes = torch.zeros((batch, head_num_q), dtype=torch.float32).to(device)
            alibi_slopes.uniform_(0, 0.125)
        softmax_scale = 1 / math.sqrt(head_size)
        hardware_time, e2e_time = benchmark_forward(tmo.single_query_cached_kv_attn,
                                                    input_q,
                                                    key_cache,
                                                    value_cache,
                                                    None,
                                                    block_tables,
                                                    context_lens,
                                                    key_cache_scale,
                                                    value_cache_scale,
                                                    alibi_slopes,
                                                    max_context_len,
                                                    -1,
                                                    -1,
                                                    softmax_scale,
                                                    repeats=args.repeat_times)
        content = [f"{batch}", f"{max_seq_len}", f"{head_num_q}", f"{head_num_kv}", f"{head_size}", f"{block_size}", f"{alibi_bias}", f"{kv_cache_dtype}", f"{use_paged_attn}", 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()
