import os

import json
import torch
import triton
import random
import itertools
import argparse
from typing import Optional
from aiter.ops.triton.chunked_pa_prefill import _kernel_paged_attention_2d
from aiter.ops.triton.chunked_pa_prefill import paged_attention_2d as paged_attention_2d_ori

_is_hip = True

version = triton.__version__.split(".")
major_version, minor_version = eval(version[0]), eval(version[1])
os.environ["AMDGCN_USE_BUFFER_OPS"] = "1"

def seed_everything(seed):
    random.seed(seed)
    torch.manual_seed(seed)

def input_helper(
    BS,
    MAX_SEQ_LEN,
    MAX_CTX_LEN,
    num_blocks,
    block_size,
    max_block_per_request,
    num_heads: int,
    head_size: int,
    num_queries_per_kv: int,
    dtype: torch.dtype,
    kv_dtype: Optional[torch.dtype],
    use_alibi_slope: bool,
    device,
):
    seed_everything(0)
    torch.set_default_device(device)

    assert block_size * max_block_per_request >= MAX_SEQ_LEN + MAX_CTX_LEN, "cache block count may not enough for each request."
    assert num_blocks >= max_block_per_request * BS, "cache block count is not enough for total requests."

    if use_alibi_slope:
        alibi_slopes = _get_alibi_slopes(num_heads).to(device)

    query_lens = [1 for _ in range(BS)]
    ctx_lens = [random.randint(16, MAX_CTX_LEN) for _ in range(BS)]
    # ctx_lens = [MAX_CTX_LEN - 1 for _ in range(BS)]
    seq_lens = [a + b for a, b in zip(query_lens, ctx_lens)]
    # print(f"{seq_lens=}")
    num_kv_heads = num_heads // num_queries_per_kv

    num_tokens = sum(query_lens)
    query = torch.randn(num_tokens, num_heads, head_size, dtype=dtype)
    output = torch.empty(num_tokens, num_heads, head_size, dtype=dtype)

    k_cache = torch.randn(
        num_blocks, block_size, num_kv_heads, head_size, 
        dtype=dtype
    )
    v_cache = torch.randn_like(k_cache)
    k_scale = None
    v_scale = None
    maybe_quantized_k_cache = k_cache
    maybe_quantized_v_cache = v_cache
    if kv_dtype is not None:
        maybe_quantized_k_cache = k_cache.to(kv_dtype)
        maybe_quantized_v_cache = v_cache.to(kv_dtype)

        scale_shape = (1,)
        k_scale = torch.rand(scale_shape, dtype=torch.float32)
        v_scale = torch.rand(scale_shape, dtype=torch.float32)

    values = torch.arange(0, num_blocks, dtype=torch.long)
    values = values[torch.randperm(num_blocks)]
    block_table = values[: BS * max_block_per_request].view(BS, max_block_per_request)
    b_seq_len = torch.tensor(seq_lens, dtype=torch.long)
    b_start_loc = torch.cumsum(torch.tensor([0] + query_lens, dtype=torch.long), dim=0)
    # transpose K_cache[num_blocks, block_size, num_kv_heads, head_size]
    # to K_cache[num_blocks, num_kv_heads, head_size/8, block_size, 8]
    k_cache = (
        maybe_quantized_k_cache.view(-1, block_size, num_kv_heads, head_size // 8, 8)
        .permute(0, 2, 3, 1, 4)
        .contiguous()
    )
    # transpose V_cache[num_blocks, block_size, num_kv_heads, head_size]
    # to V_cache[num_blocks, num_kv_heads, head_size, block_size]
    v_cache = (
        maybe_quantized_v_cache.view(-1, block_size, num_kv_heads, head_size)
        .permute(0, 2, 3, 1)
        .contiguous()
    )

    if use_alibi_slope:
        return (
            query,
            output,
            k_cache,
            v_cache,
            block_table,
            b_start_loc,
            b_seq_len,
            k_scale,
            v_scale,
            alibi_slopes,
        )
    else:
        return (
            query,
            output,
            k_cache,
            v_cache,
            block_table,
            b_start_loc,
            b_seq_len,
            k_scale,
            v_scale,
            None,
        )


def get_bench_inputs():
    names = ["BS", "MAX_SEQ_LEN", "MAX_CTX_LEN", "num_blocks", "cache_block_size", "max_block_per_request",
             "num_heads", "head_size", "num_queries_per_kv", "use_alibi_slope", "kv_dtype", "filter_by_query_len"]
    vals = []
    shapes = [
        #BS   Q_LEN  KV_LEN  BLOCKS  blocksize block_per_q  q_head head_size  q_head_per_kv
        (32,  1,     2048,   4096,    32,       128,        32,    128,       1),
        (32,  1,     2048,   4096,    32,       128,        32,    128,       4),
        (32,  1,     2048,   4096,    32,       128,        32,    128,       8),
        (32,  1,     2048,   4096,    32,       128,        32,    128,       16),
        (32,  1,     2048,   1024,    128,       32,         32,    128,       1),
        (32,  1,     2048,   1024,    128,       32,         32,    128,       4),
        (32,  1,     2048,   1024,    128,       32,         32,    128,       8),
        (32,  1,     2048,   1024,    128,       32,         32,    128,       16),
        # to add case head_size not equal head_size_padded if needed
    ]
    for use_alibi_slope in [False]:
        for kv_dtype in [None, torch.float8_e4m3fn]:
            for filter_by_query_len in [True]:
                for s in shapes:
                    vals.append((*s, use_alibi_slope, kv_dtype, filter_by_query_len))
    return names, vals


def generate_configs(config):
    keys = list(config.keys())
    values = list(config.values())
    configs_list = []
    for combination in itertools.product(*values):
        cfg = dict(zip(keys, combination))
        configs_list.append(cfg)
    return configs_list


def get_triton_configs():
    config = {
        "BLOCK_SIZE": [16, 32, 64, 128],
        "waves_per_eu": [1],
        "num_warps": [2, 4, 8],
        "instruction_sched_variant": ["none", "local-prefetch"],
        # "schedule_hint": ["none", "llvm-iglp-8", "local-prefetch"],
        "num_stages": [1, 2],
        "sched_latency": ["none", "mmac5-ds10"],
        "USE_MATRIX_LOAD": [True, False],
    }

    tt_configs = []
    for c in generate_configs(config):
        num_warps = c['num_warps']
        num_stages = c['num_stages']
        del c['num_warps']
        del c['num_stages']
        tt_configs.append(triton.Config(c, num_warps=num_warps, num_stages=num_stages))

    return tt_configs


def prune_configs(configs, nargs, **kwargs):
    def _prune(config, nargs, **kwargs):
        c = config.all_kwargs()
        num_stages = c['num_stages']
        block_size = c['BLOCK_SIZE']
        use_mls = c['USE_MATRIX_LOAD']

        all_args = {**nargs, **kwargs}
        value_cache = all_args['value_cache_ptr']
        key_cache = all_args['key_cache_ptr']
        query = all_args['query_ptr']
        cache_block_size = value_cache.shape[3]
        head_size = query.shape[2]
        head_size_padded = triton.next_power_of_2(head_size)
        cache_ele_size = key_cache.element_size()

        if use_mls:
            if block_size < 32:
                return True
            elif (cache_ele_size == 1 and block_size < 64):
                return True
            # if num_stages > 1:
            #     return True
        if cache_ele_size == 1 and block_size < 32:
            return True
        if cache_block_size % block_size != 0:
            return True
        if block_size * head_size_padded * cache_ele_size > 16384:
            return True

        return False

    res = [c for c in configs if not _prune(c, nargs, **kwargs)]
    # print(f"pruned config counts={len(res)}")
    return res

key = [
    'CACHE_BLOCK_SIZE',
    'HEAD_SIZE_PADDED',
    'num_queries_per_kv',
    'SLIDING_WINDOW',
    'USE_ALIBI_SLOPES',
    'SKIP_PREFILL',
    'HEAD_DIM_PAD_REQ'
]
fn = triton.utils.hcutune(configs=get_triton_configs(), key=key, perf_debug=True,
                          prune_configs_by={"early_config_prune": prune_configs})(_kernel_paged_attention_2d)

@torch.inference_mode()
def paged_attention_2d(
    query,
    output,
    kv_cache_dtype,
    key_cache,
    value_cache,
    block_table,
    query_start_loc,
    seq_lens,
    k_scale,
    v_scale,
    alibi_slopes=None,
    sliding_window=None,
    sm_scale=None,
    filter_by_query_len=True,
):
    if sm_scale is None:
        sm_scale = 1.0 / (query.shape[1]**0.5)
    use_alibi_slopes = alibi_slopes is not None

    if sliding_window is None or sliding_window <= 0:
        sliding_window = 0

    cache_block_size = value_cache.shape[3]
    head_size = query.shape[2]
    head_size_padded = triton.next_power_of_2(head_size)

    num_seqs = len(seq_lens)
    num_query_heads = query.shape[1]
    num_kv_heads = key_cache.shape[1]
    num_queries_per_kv = query.shape[1] // key_cache.shape[1]
    num_queries_per_kv_padded = max(triton.next_power_of_2(num_queries_per_kv), 16)

    fn[(
        num_seqs,
        num_kv_heads,
    )](
        output_ptr=output,
        query_ptr=query,
        key_cache_ptr=key_cache,
        value_cache_ptr=value_cache,
        block_tables_ptr=block_table,
        seq_lens_ptr=seq_lens,
        alibi_slopes_ptr=alibi_slopes,
        scale=sm_scale,
        k_scale=k_scale,
        v_scale=v_scale,
        num_query_heads=num_query_heads,
        num_queries_per_kv=num_queries_per_kv,
        num_queries_per_kv_padded=num_queries_per_kv_padded,
        block_table_stride=block_table.stride(0),
        query_stride_0=query.stride(0),
        query_stride_1=query.stride(1),
        output_stride_0=output.stride(0),
        output_stride_1=output.stride(1),
        CACHE_BLOCK_SIZE=cache_block_size,
        HEAD_SIZE=head_size,
        HEAD_SIZE_PADDED=head_size_padded,
        USE_ALIBI_SLOPES=use_alibi_slopes,
        SLIDING_WINDOW=sliding_window,
        x=key_cache.shape[4],
        stride_k_cache_0=key_cache.stride(0),
        stride_k_cache_1=key_cache.stride(1),
        stride_k_cache_2=key_cache.stride(2),
        stride_k_cache_3=key_cache.stride(3),
        stride_k_cache_4=key_cache.stride(4),
        stride_v_cache_0=value_cache.stride(0),
        stride_v_cache_1=value_cache.stride(1),
        stride_v_cache_2=value_cache.stride(2),
        stride_v_cache_3=value_cache.stride(3),
        SKIP_PREFILL=filter_by_query_len,
        query_start_len_ptr=query_start_loc,
        HEAD_DIM_PAD_REQ=(head_size != head_size_padded),
    )

x_names, x_vals = get_bench_inputs()
configs = [
    triton.testing.Benchmark(
        x_names=x_names,
        x_vals=x_vals,
        line_arg="provider",
        line_vals=["triton"],
        line_names=["triton"],
        styles=[("red", "-")],
        ylabel="ms",
        plot_name="paged_attention_2d",
        args={'dtype': torch.float16},
    )
]

@triton.utils.dist_perf_report(configs)
def bench_paged_attention_2d(BS, MAX_SEQ_LEN, MAX_CTX_LEN, num_blocks, cache_block_size, max_block_per_request,
                            num_heads, head_size, num_queries_per_kv, use_alibi_slope, kv_dtype, 
                            filter_by_query_len, provider, dtype):
    device = 'cpu' if os.getenv("TRITON_HCUTUNE_COMPILE_ONLY", "") == "1" else "cuda"

    (
        query,
        output,
        k_cache,
        v_cache,
        block_table,
        b_start_loc,
        b_seq_len,
        k_scale,
        v_scale,
        alibi_slopes,
    ) = input_helper(
        BS,
        MAX_SEQ_LEN,
        MAX_CTX_LEN,
        num_blocks,
        cache_block_size,
        max_block_per_request,
        num_heads,
        head_size,
        num_queries_per_kv,
        dtype,
        kv_dtype,
        use_alibi_slope,
        device
    )
    tuning_mod = int(os.environ.get("TUNING_MOD", "1"))
    if tuning_mod:
        fn = lambda: paged_attention_2d(
            query,
            output,
            "auto",
            k_cache,
            v_cache,
            block_table,
            b_start_loc,
            b_seq_len,
            k_scale,
            v_scale,
            alibi_slopes=alibi_slopes,
            filter_by_query_len=filter_by_query_len,
        )
    else:
        fn = lambda: paged_attention_2d_ori(
            query,
            output,
            "auto",
            k_cache,
            v_cache,
            block_table,
            b_start_loc,
            b_seq_len,
            k_scale,
            v_scale,
            alibi_slopes=alibi_slopes,
            filter_by_query_len=filter_by_query_len,
        )
    return triton.testing.do_bench(fn)

def parse_args():
    parser = argparse.ArgumentParser()

    parser.add_argument("--perf", action='store_true', default=False,
                        help='benchmark with hcutuner perf mode')

    return parser.parse_args()


if __name__ == "__main__":
    args = parse_args()
    if args.perf:
        os.environ["TRITON_HCUTUNE_PERF_MODE"] = "1"

    bench_paged_attention_2d.run(print_data=True, save_path=f"./tune_paged_attention_2d_out")
