import argparse
import random
import os

import torch
import torch_mlu
import torch_mlu_ops as tmo

from common import benchmark_forward, save_to_csv
from itertools import product
from tabulate import tabulate

e2e_time_param_dict_list = [
    {"max_batch_size": 128, "batch_size": [1, 32, 64], "max_context_len": [1024, 2048, 3072, 4096],
     "head_num_q": 32, "head_num_kv": 1, "cache_mem_len": 6144, "head_size": 128,
     "input_dtype": [torch.float16, torch.bfloat16], "quant_mode": [0, 1], "quant_bit": [4, 8],
     "use_offset": True},
]

def main():
    parser = argparse.ArgumentParser(description="Benchmark for dequant from linear cache.")
    parser.add_argument('--repeat_times', type=int, default=100, 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 = ["max_batch_size", "batch_size", "max_context_len", "head_num_q", "head_num_kv",
              "cache_mem_len", "head_size", "input_dytpe", "quant_mode", "quant_bit",
              "use_offset", "hardware_time(us)", "e2e_latency(us)"]
    contents = []
    mlu_name = torch.mlu.get_device_name()
    for params_dict in e2e_time_param_dict_list:
        max_batch_size = params_dict["max_batch_size"]
        batch_size_list = params_dict["batch_size"]
        max_context_len_list = params_dict["max_context_len"]
        head_num_q = params_dict["head_num_q"]
        head_num_kv = params_dict["head_num_kv"]
        head_size = params_dict["head_size"]
        cache_mem_len = params_dict["cache_mem_len"]
        input_dtype_list = params_dict["input_dtype"]
        quant_mode_list = params_dict["quant_mode"]
        quant_bit_list = params_dict["quant_bit"]
        use_offset = params_dict["use_offset"]
        for batch_size, max_context_len, quant_mode, quant_bit, dtype in list(product( \
                batch_size_list, max_context_len_list, quant_mode_list, quant_bit_list, \
                input_dtype_list)):
            torch.manual_seed(2766)
            torch.mlu.manual_seed(2766)
            if dtype == torch.bfloat16 and not torch_mlu.mlu.is_bf16_supported():
                continue
            if "MLU3" in mlu_name and (2 * cache_mem_len * max_batch_size * head_num_kv \
                * head_size >= 2**31 - 1):
                print("large tensor is not support on {}, skip".format(mlu_name))
                continue
            total_heads = head_num_q + head_num_kv * 2
            assert max_context_len <= cache_mem_len, "max_context_len should smaller than or " \
                "equal to cache_mem_len."
            max_seq_offset = cache_mem_len - max_context_len
            # Generates key and cache from context
            context_lens = torch.randint(size=[batch_size], low=max_context_len,
                                         high=max_context_len + 1,
                                         dtype=torch.int32, device="mlu")
            if use_offset:
                context_paddings = torch.randint(size=[batch_size], low=0, high=max_seq_offset,
                                                 dtype=torch.int32, device="mlu")
            else:
                context_paddings = torch.zeros_like(context_lens)
            cu_context_lens = torch.cumsum(context_lens + context_paddings, dim=-1)
            total_seqlen = cu_context_lens[-1]
            context_seq_offset = torch.zeros([batch_size], dtype=torch.int32, device="mlu")
            context_seq_offset[1:] = cu_context_lens[:-1]
            context = torch.randn([total_seqlen, total_heads, head_size], dtype=dtype, device="mlu")
            key = context[..., head_num_q:head_num_q + head_num_kv, :]
            value = context[..., head_num_q + head_num_kv:head_num_q + 2 * head_num_kv, :]

            # Generates key_cache and value_cache
            cache_bs_id = torch.IntTensor(random.sample([*range(0, batch_size + 1)], batch_size)).mlu()
            cache_seq_offset = torch.randint(low=-1, high=max_seq_offset, size=[batch_size],
                                             dtype=torch.int32, device="mlu")
            if quant_bit == 4:
                key_cache = torch.randint(low=-128, high=127, dtype=torch.int32, size=(max_batch_size,
                                          head_num_kv, cache_mem_len, head_size // 2), device="mlu")
                value_cache = torch.randint(low=-128, high=127, dtype=torch.int32, size=(max_batch_size,
                                            head_num_kv, cache_mem_len // 2, head_size), device="mlu")
                key_cache, value_cache = key_cache.to(torch.int8), value_cache.to(torch.int8)
            else:
                cache = torch.randint(size=(2, max_batch_size, head_num_kv, cache_mem_len, head_size),
                                      low=-128, high=127, dtype=torch.int32, device="mlu")
                cache = cache.to(torch.int8)
                key_cache, value_cache = cache[[0, 1]]

            # Generates key_cache_scale and value_cache_scale
            if quant_mode == 0:  # quant_mode == 0 is per channel
                cache_scale = torch.randn((2, head_num_kv, head_size), dtype=torch.float, device="mlu")
            else:  # quant_mode != 1 (== 1 for extend) is per head
                cache_scale = torch.randn((2, max_batch_size, head_num_kv, cache_mem_len),
                                           dtype=torch.float, device="mlu")
            key_cache_scale, value_cache_scale = cache_scale[[0, 1]]
            hardware_time, e2e_time = benchmark_forward(tmo.dequant_from_linear_cache,
                                                        key, value, key_cache, value_cache,
                                                        key_cache_scale, value_cache_scale,
                                                        context_lens, max_context_len,
                                                        context_seq_offset if use_offset else None,
                                                        cache_bs_id, cache_seq_offset, quant_mode,
                                                        quant_bit, repeats=args.repeat_times)
            content = [f"{max_batch_size}", f"{batch_size}", f"{max_context_len}", f"{head_num_q}",
                       f"{head_num_kv}", f"{cache_mem_len}", f"{head_size}", f"{dtype}", f"{quant_mode}",
                       f"{quant_bit}", f"{quant_mode}", f"{use_offset}", 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()
