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

e2e_time_param_dict_list = [{"batch": 16, "max_seq_len_lp": 4064, "max_seq_len_hp": 32, "head_num_q": 8, "head_num_kv": 1,
                              "head_size": 128, "alibi_bias": False, "quant_bit_lp": 4, "quant_bit_hp": 8,
                              "use_paged_attn": False, "input_dtype": [torch.float16, torch.bfloat16]},
                            {"batch": 128, "max_seq_len_lp": 4064, "max_seq_len_hp": 32, "head_num_q": 8, "head_num_kv": 1,
                              "head_size": 128, "alibi_bias": False, "quant_bit_lp": 4, "quant_bit_hp": 8,
                              "use_paged_attn": False, "input_dtype": [torch.float16, torch.bfloat16]},
                            {"batch": 512, "max_seq_len_lp": 4064, "max_seq_len_hp": 32, "head_num_q": 8, "head_num_kv": 1,
                              "head_size": 128, "alibi_bias": True, "quant_bit_lp": 4, "quant_bit_hp": 8,
                              "use_paged_attn": False, "input_dtype": [torch.float16, torch.bfloat16]},
                            {"batch": 16, "max_seq_len_lp": 32736, "max_seq_len_hp": 32, "head_num_q": 8, "head_num_kv": 1,
                              "head_size": 128, "alibi_bias": False, "quant_bit_lp": 4, "quant_bit_hp": 8,
                              "use_paged_attn": False, "input_dtype": [torch.float16, torch.bfloat16]},
                            {"batch": 128, "max_seq_len_lp": 32736, "max_seq_len_hp": 32, "head_num_q": 8, "head_num_kv": 1,
                              "head_size": 128, "alibi_bias": False, "quant_bit_lp": 4, "quant_bit_hp": 8,
                              "use_paged_attn": False, "input_dtype": [torch.float16, torch.bfloat16]}]

def gen_cache(batch, num_kv_heads, head_size, is_pagedattn, max_context_len, data_type, quant_bit, quant_mode):
    int_max = float(2 ** (quant_bit - 1) - 1)
    int_min = -float(2 ** (quant_bit - 1))
    context_lens = torch.randint(max_context_len, max_context_len + 1, (batch, ), dtype=torch.int32).mlu()
    block_size = 16
    if is_pagedattn is False:
        block_size = max_context_len
    num_blocks = (int)(batch * ((max_context_len + block_size - 1)/ block_size))
    max_num_blocks_per_seq = (max_context_len + block_size - 1) // block_size
    cache_shape = (num_blocks, num_kv_heads, block_size, head_size)
    if quant_mode == "per_token":
        scale_shape = (num_blocks, num_kv_heads, block_size, 1)
    else: # per channel
        scale_shape = (num_kv_heads, head_size)

    if quant_bit == 4:
        cache_shape_k_int4 = (num_blocks, num_kv_heads, block_size, head_size//2)
        cache_shape_v_int4 = (num_blocks, num_kv_heads, block_size//2, head_size)
        key_cache = torch.zeros(cache_shape_k_int4).uniform_(int_min, int_max).to(torch.int8).mlu()
        value_cache = torch.zeros(cache_shape_v_int4).uniform_(int_min, int_max).to(torch.int8).mlu()
        key_cache_scale = torch.randn(size=scale_shape, dtype=torch.float32).mlu()
        value_cache_scale = torch.randn(size=scale_shape, dtype=torch.float32).mlu()
    elif quant_bit == 8:
        key_cache = torch.zeros(cache_shape).uniform_(int_min, int_max).to(torch.int8).mlu()
        value_cache = torch.zeros(cache_shape).uniform_(int_min, int_max).to(torch.int8).mlu()
        key_cache_scale = torch.randn(size=scale_shape, dtype=torch.float32).mlu()
        value_cache_scale = torch.randn(size=scale_shape, dtype=torch.float32).mlu()
    elif quant_bit == -1:
        key_cache = torch.randn(cache_shape, dtype=data_type).mlu()
        value_cache = torch.randn(cache_shape, dtype=data_type).mlu()
        key_cache_scale = None
        value_cache_scale = None
    else:
        print("!!!!!!!!!!!gen case error, quant_bit must be in {-1, 4, 8}")
    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)
    return key_cache, value_cache, key_cache_scale, value_cache_scale, context_lens, block_tables


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_lp", "max_seq_len_hp", "head_num_q", "head_num_kv", "head_size", "alibi_bias", "quant_bit_lp", "quant_bit_hp","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_lp = params_dict["max_seq_len_lp"]
    max_seq_len_hp = params_dict["max_seq_len_hp"]
    head_num_q = params_dict["head_num_q"]
    head_num_kv = params_dict["head_num_kv"]
    head_size = params_dict["head_size"]
    alibi_bias = params_dict["alibi_bias"]
    quant_bit_lp = params_dict["quant_bit_lp"]
    quant_bit_hp = params_dict["quant_bit_hp"]
    use_paged_attn = params_dict["use_paged_attn"]
    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_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, :]

        params_lp = gen_cache(batch, head_num_kv, head_size, use_paged_attn, max_seq_len_lp, dtype, quant_bit_lp, "per_token")
        params_hp = gen_cache(batch, head_num_kv, head_size, use_paged_attn, max_seq_len_hp, dtype, quant_bit_hp, "per_channel")
        key_cache_lp, value_cache_lp, key_cache_scale_lp, value_cache_scale_lp, context_lens_lp, block_tables_lp = params_lp
        key_cache_hp, value_cache_hp, key_cache_scale_hp, value_cache_scale_hp, context_lens_hp, block_tables_hp = params_hp

        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_mixed_cached_kv_attn,
                                                    input_q,
                                                    key_cache_lp, value_cache_lp,
                                                    key_cache_hp, value_cache_hp,
                                                    None, #output
                                                    block_tables_lp, block_tables_hp,
                                                    context_lens_lp, context_lens_hp,
                                                    key_cache_scale_lp, value_cache_scale_lp,
                                                    key_cache_scale_hp, value_cache_scale_hp,
                                                    alibi_slopes,
                                                    max_seq_len_lp, max_seq_len_hp,
                                                    softmax_scale, True,
                                                    quant_bit_lp, quant_bit_hp,
                                                    repeats=args.repeat_times)
        content = [f"{batch}", f"{max_seq_len_lp}",  f"{max_seq_len_hp}", f"{head_num_q}", f"{head_num_kv}", f"{head_size}", f"{alibi_bias}", f"{quant_bit_lp}", f"{quant_bit_hp}", 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()
