import os
import json
import torch
import triton
import random
import itertools

from aiter.ops.triton.grouped_decode_attention import _fwd_grouped_kernel_stage1, _fwd_kernel_stage2


def input_helper(B, S, H_Q, H_KV, D, D_V, max_kv_splits, dtype, device='cuda'):
    torch.manual_seed(0)

    seq_len = S  # This represents the number of tokens already in the sequence
    total_tokens = B * seq_len
    sm_scale = 1.0 / (D**0.5)
    num_kv_splits = torch.full((B,), 4, dtype=torch.int32, device=device)

    # q represents the new token being generated, one per batch
    q = torch.randn(B, H_Q, D, dtype=dtype, device=device)

    # k_buffer and v_buffer represent all previous tokens
    k_buffer = torch.randn(total_tokens, H_KV, D, dtype=dtype, device=device)
    v_buffer = torch.randn(total_tokens, H_KV, D_V, dtype=dtype, device=device)

    o_grouped = torch.zeros(B, H_Q, D_V, dtype=dtype, device=device)

    b_seq_len = torch.full((B,), seq_len, device=device)

    kv_indptr = torch.zeros((B + 1,), dtype=torch.int32, device=device)
    kv_indptr[1 : B + 1] = torch.cumsum(b_seq_len[:B], dim=0)
    kv_indices = torch.arange(total_tokens, dtype=torch.int32, device=device)

    attn_logits = torch.empty(
        (B, H_Q, max_kv_splits, D_V),
        dtype=torch.float32,
        device=device,
    )
    attn_lse = torch.empty(
        (B, H_Q, max_kv_splits),
        dtype=torch.float32,
        device=device,
    )

    return (
        q,
        k_buffer,
        v_buffer,
        o_grouped,
        kv_indptr,
        kv_indices,
        attn_logits,
        attn_lse,
        num_kv_splits,
        sm_scale,
    )


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_stage1_triton_configs():
    config = {
        "BLOCK_N": [16, 32, 64],
        "BLOCK_H": [16, 32, 64],
        "waves_per_eu": [1],
        "num_warps": [4, 8, 16],
        "matrix_instr_nonkdim": [16],
        # "instruction_sched_variant": ["none", "llvm-iglp-1", "llvm-iglp-8", "local-prefetch"],
        "num_stages": [1, 2, 3],
        "sched_latency": ["none", "mmac5-ds10"],
        "kpack": [1, 2],
    }
    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 get_stage2_triton_configs():
    config = {
        "waves_per_eu": [1, 2, 4, 8],
        "num_warps": [2, 4, 8, 16],
        "matrix_instr_nonkdim": [16],
        # "instruction_sched_variant": ["none", "llvm-iglp-1", "llvm-iglp-8", "local-prefetch"],
        "num_stages": [1, 2, 3],
        "sched_latency": ["none", "mmac5-ds10"],
        "kpack": [1, 2],
    }
    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):
#         c = config.all_kwargs()
#         if (c['BLOCK_H'] >= 32 and c['BLOCK_N'] >= 32) and c['num_stages'] >= 2:
#             return True

#     res = [c for c in configs if not _prune(c)]
#     return res


key = [
    'kv_group_num',
    'q_head_num',
    'Lk',
    'Lv',
]
# fn_stage1 = triton.utils.hcutune(configs=get_stage1_triton_configs(), key=key, perf_debug=True, warmup=100, rep=300, prune_configs_by={"early_config_prune": prune_configs})(_fwd_grouped_kernel_stage1)
fn_stage1 = triton.utils.hcutune(configs=get_stage1_triton_configs(), key=key,
                                 perf_debug=True)(_fwd_grouped_kernel_stage1)


key = [
    'MAX_KV_SPLITS',
    'Lv',
]
fn_stage2 = triton.utils.hcutune(configs=get_stage2_triton_configs(), key=key,
                                 perf_debug=True, warmup=100, rep=300)(_fwd_kernel_stage2)


_MIN_BLOCK_KV = 32


def _decode_softmax_reducev_fwd(
    logits,
    lse,
    q,
    o,
    v_buffer,
    kv_indptr,
    num_kv_splits,
    max_kv_splits,
):
    batch, head_num = q.shape[0], q.shape[1]
    Lv = v_buffer.shape[-1]
    BLOCK_DV = triton.next_power_of_2(Lv)

    MAX_KV_SPLITS = max_kv_splits

#    extra_kargs = {}
#    if _is_hip:
#        extra_kargs = {"waves_per_eu": 4, "matrix_instr_nonkdim": 16, "kpack": 2}

    grid = (batch, head_num)
    fn_stage2[grid](
        logits,
        lse,
        o,
        kv_indptr,
        num_kv_splits,
        logits.stride(0),
        logits.stride(1),
        logits.stride(2),
        o.stride(0),
        o.stride(1),
        MAX_KV_SPLITS=MAX_KV_SPLITS,
        MIN_BLOCK_KV=_MIN_BLOCK_KV,
        BLOCK_DV=BLOCK_DV,
        Lv=Lv,
#        num_warps=4,
#        num_stages=2,
#        **extra_kargs,
    )


def _decode_grouped_att_m_fwd(
    q,
    k_buffer,
    v_buffer,
    att_out,
    att_lse,
    kv_indptr,
    kv_indices,
    num_kv_splits,
    max_kv_splits,
    sm_scale,
    logit_cap,
):
#    BLOCK = 32
    Lk = k_buffer.shape[-1]
    Lv = v_buffer.shape[-1]

#    # [TODO] work around shmem limit on MI3xx
#    if _is_hip and Lk >= 576:
#        BLOCK = 16

    if Lk == 576:
        BLOCK_DMODEL = 512
        BLOCK_DPE = 64
    elif Lk == 288:
        BLOCK_DMODEL = 256
        BLOCK_DPE = 32
    else:
        BLOCK_DMODEL = triton.next_power_of_2(Lk)
        BLOCK_DPE = 0
    BLOCK_DV = triton.next_power_of_2(Lv)

    batch, head_num = kv_indptr.shape[0] - 1, q.shape[1]
    kv_group_num = q.shape[1] // k_buffer.shape[1]

    BLOCK_H = 16
    MAX_KV_SPLITS = max_kv_splits
    grid = lambda META: (
        batch,
        triton.cdiv(head_num, min(META['BLOCK_H'], kv_group_num)),
        MAX_KV_SPLITS,
    )

#    extra_kargs = {}
#    num_stages = 2
#    if _is_hip:
#        extra_kargs = {"waves_per_eu": 1, "matrix_instr_nonkdim": 16, "kpack": 2}
#        num_stages = 1

    fn_stage1[grid](
        q,
        k_buffer,
        v_buffer,
        sm_scale,
        kv_indptr,
        kv_indices,
        att_out,
        att_lse,
        num_kv_splits,
        q.stride(0),
        q.stride(1),
        k_buffer.stride(0),
        k_buffer.stride(1),
        v_buffer.stride(0),
        v_buffer.stride(1),
        att_out.stride(0),
        att_out.stride(1),
        att_out.stride(2),
        kv_group_num=kv_group_num,
        q_head_num=head_num,
        BLOCK_DMODEL=BLOCK_DMODEL,
        BLOCK_DPE=BLOCK_DPE,
        BLOCK_DV=BLOCK_DV,
#        BLOCK_N=BLOCK,
#        BLOCK_H=BLOCK_H,
        MIN_BLOCK_KV=_MIN_BLOCK_KV,
        logit_cap=logit_cap,
#        num_warps=4,
#        num_stages=num_stages,
        Lk=Lk,
        Lv=Lv,
#        **extra_kargs,
    )


def decode_attention_fwd_grouped(
    q,
    k_buffer,
    v_buffer,
    o,
    kv_indptr,
    kv_indices,
    attn_logits,
    attn_lse,
    num_kv_splits,
    max_kv_splits,
    sm_scale,
    logit_cap=0.0,
):
    _decode_grouped_att_m_fwd(
        q,
        k_buffer,
        v_buffer,
        attn_logits,
        attn_lse,
        kv_indptr,
        kv_indices,
        num_kv_splits,
        max_kv_splits,
        sm_scale,
        logit_cap,
    )
    _decode_softmax_reducev_fwd(
        attn_logits,
        attn_lse,
        q,
        o,
        v_buffer,
        kv_indptr,
        num_kv_splits,
        max_kv_splits,
    )


def get_bench_inputs():
    names = ["B", "S", "H_Q", "H_KV", "D", "D_V", "max_kv_splits"]
    vals = []
    shapes = [
        (1, 8, 16, 1, 576, 512),
        (1, 13, 16, 1, 576, 512),
        (1, 135, 16, 1, 576, 512),
        (4, 152, 16, 1, 576, 512),
        (4, 256, 16, 1, 576, 512),
        (4, 593, 16, 1, 576, 512),
    ]

    for max_kv_splits in [16]:
        for s in shapes:
            vals.append((*s, max_kv_splits))

    return names, vals


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="extend_attention",
        args={"dtype": torch.float16},
    )
]


@triton.utils.dist_perf_report(configs)
def bench_decode_attention(B, S, H_Q, H_KV, D, D_V, max_kv_splits, provider, dtype):
    device = 'cpu' if os.getenv("TRITON_HCUTUNE_COMPILE_ONLY", "") == "1" else "cuda"

    (
        q,
        k_buffer,
        v_buffer,
        o_grouped,
        kv_indptr,
        kv_indices,
        attn_logits,
        attn_lse,
        num_kv_splits,
        sm_scale
    ) = input_helper(B, S, H_Q, H_KV, D, D_V, max_kv_splits, dtype)

    fn = lambda: decode_attention_fwd_grouped(
        q,
        k_buffer,
        v_buffer,
        o_grouped,
        kv_indptr,
        kv_indices,
        attn_logits,
        attn_lse,
        num_kv_splits,
        max_kv_splits,
        sm_scale=sm_scale,
    )
    return triton.testing.do_bench(fn)


if __name__ == "__main__":
    bench_decode_attention.run(print_data=True)
