import os

import json
import torch
import triton
import random
import itertools
import argparse
from typing import Optional
from aiter.ops.triton.pa_prefill import _fwd_kernel, _fwd_kernel_alibi
from aiter.ops.triton.pa_prefill import context_attention_fwd as context_attention_fwd_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 = [random.randint(16, MAX_SEQ_LEN) for _ in range(BS)]
    ctx_lens = [random.randint(16, MAX_CTX_LEN) for _ in range(BS)]
    seq_lens = [a + b for a, b in zip(query_lens, ctx_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 = torch.randn(sum(query_lens), num_kv_heads, head_size, dtype=dtype)
    v = torch.randn(sum(query_lens), num_kv_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,
            k,
            v,
            output,
            k_cache,
            v_cache,
            block_table,
            b_start_loc,
            b_seq_len,
            k_scale,
            v_scale,
            alibi_slopes,
        )
    else:
        return (
            query,
            k,
            v,
            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", "skip_decode"]
    vals = []
    shapes = [
        #BS   Q_LEN  KV_LEN  BLOCKS  blocksize block_per_q  q_head head_size  q_head_per_kv
        (16,  128,  2048,   2048,    32,       128,        32,    128,       4),
        (16,  256,  2048,   2048,    32,       128,        32,    128,       4),
        (16,  512,  2048,   2048,    32,       128,        32,    128,       4),
        (16,  1024,  2048,   2048,    32,       128,        32,    128,       4),
        (16,  2048,  2048,   2048,    32,       128,        32,    128,       4),
        (16,  128,  2048,   512,    128,       32,         32,    128,        4),
        (16,  256,  2048,   512,    128,       32,         32,    128,        4),
        (16,  512,  2048,   512,    128,       32,         32,    128,        4),
        (16,  1024,  2048,   512,    128,       32,         32,    128,       4),
        (16,  2048,  2048,   512,    128,       32,         32,    128,       4),
        # 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 kv_dtype in [None]:
            for skip_decode in [True]:
                for s in shapes:
                    vals.append((*s, use_alibi_slope, kv_dtype, skip_decode))
    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(is_alibi=False):
    config = {
        "BLOCK_M": [32, 64, 128, 256],
        "BLOCK_N": [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": [False] if is_alibi else [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_n = c['BLOCK_N']
        block_m = c['BLOCK_M']
        use_mls = c['USE_MATRIX_LOAD']

        all_args = {**nargs, **kwargs}
        value_cache = all_args['V_cache']
        key_cache = all_args['K_cache']
        query = all_args['Q']
        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 block_m < block_n:
            # blockM can not less than blockN
            return True
        if use_mls:
            if block_n < 32:
                return True
            elif (cache_ele_size == 1 and block_n < 64):
                return True
            # if num_stages > 1:
            #     return True
            if cache_block_size % block_n != 0:
                return True
        if cache_ele_size == 1 and block_n < 32:
            return True
        if block_n * 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 = [
    'block_size',  #cache_block_size
    'BLOCK_DMODEL_PADDED',   #head_size_padded
    'SLIDING_WINDOW',
    'SKIP_DECODE',
    'HEAD_DIM_PAD_REQ',
    'max_input_len',
]
fn = triton.utils.hcutune(configs=get_triton_configs(), key=key, perf_debug=True,
                          prune_configs_by={"early_config_prune": prune_configs})(_fwd_kernel)
fn_alibi = triton.utils.hcutune(configs=get_triton_configs(True), key=key, perf_debug=True,
                          prune_configs_by={"early_config_prune": prune_configs})(_fwd_kernel_alibi)


@torch.inference_mode()
def context_attention_fwd(q,
                          k,
                          v,
                          o,
                          kv_cache_dtype: str,
                          k_cache,
                          v_cache,
                          b_loc,
                          b_start_loc,
                          b_seq_len,
                          max_input_len,
                          k_scale: torch.Tensor,
                          v_scale: torch.Tensor,
                          alibi_slopes=None,
                          sliding_window=None,
                          sm_scale=None,
                          skip_decode=False):

    q_dtype_is_f32 = q.dtype is torch.float32
    # shape constraints
    Lq, Lk, Lv = q.shape[-1], k.shape[-1], v.shape[-1]
    assert Lq == Lk and Lk == Lv
    # round up Lk to a power of 2 - this is required for Triton block size
    Lk_padded = triton.next_power_of_2(Lk)

    # Turing does have tensor core for float32 multiplication
    # use ieee as fallback for triton kernels work. There is also
    # warning on vllm/config.py to inform users this fallback
    # implementation
    IN_PRECISION = None
    if sm_scale is None:
        sm_scale = 1.0 / (Lq**0.5)
    batch, head = b_seq_len.shape[0], q.shape[1]
    num_queries_per_kv = q.shape[1] // k.shape[1]

    assert batch + 1 == len(b_start_loc)
    grid = lambda META: (batch, head, triton.cdiv(max_input_len, META['BLOCK_M']))  # batch, head,

    # 0 means "disable"
    if sliding_window is None or sliding_window <= 0:
        sliding_window = 0

    if "fp8" in kv_cache_dtype and (k_cache.dtype == torch.uint8 or v_cache.dtype == torch.uint8):
        # kv_cache may view as uint8
        if kv_cache_dtype in ("fp8", "fp8e4m3"):
            target_dtype = torch.float8_e4m3fn
        elif kv_cache_dtype == "fp8e5m2":
            target_dtype = torch.float8_e5m2
        else:
            raise ValueError("Unsupported FP8 dtype:", kv_cache_dtype)
        k_cache = k_cache.view(target_dtype)
        v_cache = v_cache.view(target_dtype)

    if alibi_slopes is not None:
        fn_alibi[grid](
            q,
            k,
            v,
            k_cache,
            v_cache,
            b_loc,
            sm_scale,
            k_scale,
            v_scale,
            b_start_loc,
            b_seq_len,
            alibi_slopes,
            v_cache.shape[3],
            k_cache.shape[4],
            o,
            b_loc.stride(0),
            b_loc.stride(1),
            q.stride(0),
            q.stride(1),
            q.stride(2),
            k.stride(0),
            k.stride(1),
            k.stride(2),
            v.stride(0),
            v.stride(1),
            v.stride(2),
            o.stride(0),
            o.stride(1),
            o.stride(2),
            k_cache.stride(0),
            k_cache.stride(1),
            k_cache.stride(2),
            k_cache.stride(3),
            k_cache.stride(
                4
            ),  #[num_blocks, num_kv_heads, head_size/x, block_size, x]
            v_cache.stride(0),
            v_cache.stride(1),
            v_cache.stride(2),
            v_cache.stride(
                3),  #[num_blocks, num_kv_heads, head_size, block_size]
            num_queries_per_kv=num_queries_per_kv,
            IN_PRECISION=IN_PRECISION,
            # BLOCK_M=BLOCK_M,
            BLOCK_DMODEL=Lk,
            BLOCK_DMODEL_PADDED=Lk_padded,
            # BLOCK_N=BLOCK_N,
            SKIP_DECODE=skip_decode,
        )
        return

    fn[grid](
        q,
        k,
        v,
        k_cache,
        v_cache,
        b_loc,
        sm_scale,
        k_scale,
        v_scale,
        b_start_loc,
        b_seq_len,
        v_cache.shape[3],
        k_cache.shape[4],
        o,
        b_loc.stride(0),
        b_loc.stride(1),
        q.stride(0),
        q.stride(1),
        q.stride(2),
        k.stride(0),
        k.stride(1),
        k.stride(2),
        v.stride(0),
        v.stride(1),
        v.stride(2),
        o.stride(0),
        o.stride(1),
        o.stride(2),
        k_cache.stride(0),
        k_cache.stride(1),
        k_cache.stride(2),
        k_cache.stride(3),
        k_cache.stride(
            4),  #[num_blocks, num_kv_heads, head_size/x, block_size, x]
        v_cache.stride(0),
        v_cache.stride(1),
        v_cache.stride(2),
        v_cache.stride(
            3),  #[num_blocks, num_kv_heads, head_size, block_size]
        num_queries_per_kv=num_queries_per_kv,
        IN_PRECISION=IN_PRECISION,
        # BLOCK_M=BLOCK_M,
        BLOCK_DMODEL=Lk,
        BLOCK_DMODEL_PADDED=Lk_padded,
        # BLOCK_N=BLOCK_N,
        SLIDING_WINDOW=sliding_window,
        SKIP_DECODE=skip_decode,
        # USE_MATRIX_LOAD=use_matrix_load,
        HEAD_DIM_PAD_REQ=(Lk != Lk_padded),
        max_input_len=max_input_len,
    )
    return

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_context_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, skip_decode, provider, dtype):
    device = 'cpu' if os.getenv("TRITON_HCUTUNE_COMPILE_ONLY", "") == "1" else "cuda"

    (
        query,
        key,
        value,
        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: context_attention_fwd(
            query,
            key,
            value,
            output,
            "auto",
            k_cache,
            v_cache,
            block_table,
            b_start_loc,
            b_seq_len,
            MAX_SEQ_LEN,
            k_scale,
            v_scale,
            alibi_slopes=alibi_slopes,
            skip_decode=skip_decode,
        )
    else:
        fn = lambda: context_attention_fwd_ori(
            query,
            key,
            value,
            output,
            "auto",
            k_cache,
            v_cache,
            block_table,
            b_start_loc,
            b_seq_len,
            MAX_SEQ_LEN,
            k_scale,
            v_scale,
            alibi_slopes=alibi_slopes,
            skip_decode=skip_decode,
        )
    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_context_attention_2d.run(print_data=True, save_path=f"./tune_context_attention_2d_out")
