import os
import torch
import triton
import itertools
from typing import Any, Optional

from aiter.ops.triton.sage_attention_quant_per_block import quant_per_block_int8_kernel
from aiter.ops.triton.sage_attention_qk_int8_per_block import _attn_fwd
from aiter.ops.triton.sage_attention_qk_int8_per_block_causal import _attn_causal_fwd


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_space):
    tt_configs = []
    for c in generate_configs(config_space):
        if 'num_warps' in c:
            num_warps = c['num_warps']
            del c['num_warps']
        else:
            num_warps = 4

        if 'num_stages' in c:
            num_stages = c['num_stages']
            del c['num_stages']
        else:
            num_stages = 1

        if num_stages == 1 and 'instruction_sched_variant' in c and c['instruction_sched_variant'] != 'none':
            continue

        tt_configs.append(triton.Config(c, num_warps=num_warps, num_stages=num_stages))
    return tt_configs


key_causal = [
    'qo_len',
    'kv_len',
    'H',
    'num_kv_groups',
    'BLOCK_M',
    'BLOCK_N',
]
configs_causal = get_triton_configs({
    "STAGE": [1],
    "waves_per_eu": [1, 2],
    "num_warps": [4, 8, 16],
    # "instruction_sched_variant": ["none", "llvm-iglp-1", "llvm-iglp-8", "local-prefetch"],
    "matrix_instr_nonkdim": [16],
    "num_stages": [1, 2, 3],
    "sched_latency": ["none", "mmac5-ds10"],
    "kpack": [1, 2],
})
fn_causal = triton.utils.hcutune(configs=configs_causal, key=key_causal, perf_debug=True)(_attn_causal_fwd)

def attn_true(q, k, v, block_m, block_n, q_scale, k_scale, tensor_layout="HND", output_dtype=torch.float16, return_lse=False):
    o = torch.empty(q.shape, dtype=output_dtype, device=q.device)

    if tensor_layout == "HND":
        b, h_qo, qo_len, head_dim = q.shape
        _, h_kv, kv_len, _ = k.shape

        stride_bz_q, stride_h_q, stride_seq_q = q.stride(0), q.stride(1), q.stride(2)
        stride_bz_k, stride_h_k, stride_seq_k = k.stride(0), k.stride(1), k.stride(2)
        stride_bz_v, stride_h_v, stride_seq_v = v.stride(0), v.stride(1), v.stride(2)
        stride_bz_o, stride_h_o, stride_seq_o = o.stride(0), o.stride(1), o.stride(2)
    elif tensor_layout == "NHD":
        b, qo_len, h_qo, head_dim = q.shape
        _, kv_len, h_kv, _ = k.shape

        stride_bz_q, stride_h_q, stride_seq_q = q.stride(0), q.stride(2), q.stride(1)
        stride_bz_k, stride_h_k, stride_seq_k = k.stride(0), k.stride(2), k.stride(1)
        stride_bz_v, stride_h_v, stride_seq_v = v.stride(0), v.stride(2), v.stride(1)
        stride_bz_o, stride_h_o, stride_seq_o = o.stride(0), o.stride(2), o.stride(1)
    else:
        raise ValueError(f"tensor_layout {tensor_layout} not supported")
    
    assert qo_len == kv_len, "qo_len and kv_len must be equal for causal attention"

    HEAD_DIM_K = head_dim
    num_kv_groups = h_qo // h_kv

    if return_lse:
        lse = torch.empty([b, h_qo, qo_len], dtype=torch.float32, device=q.device)
    else:
        lse = torch.empty([0], dtype=torch.float32, device='cpu')

    grid = lambda META: (triton.cdiv(qo_len, block_m), h_qo, b   )
    fn_causal[grid](
        q, k, v, q_scale, k_scale, o, lse,
        stride_bz_q, stride_h_q, stride_seq_q, 
        stride_bz_k, stride_h_k, stride_seq_k,  
        stride_bz_v, stride_h_v, stride_seq_v,  
        stride_bz_o, stride_h_o, stride_seq_o,
        qo_len, kv_len,
        h_qo, num_kv_groups,
        HEAD_DIM=HEAD_DIM_K,
        RETURN_LSE=return_lse,
        BLOCK_M=block_m,
        BLOCK_N=block_n,
    )

    return o, lse


key = [
    'qo_len',
    'kv_len',
    'H',
    'num_kv_groups',
    'BLOCK_M',
    'BLOCK_N',
]
configs = get_triton_configs({
    "STAGE": [1],
    "waves_per_eu": [1, 2],
    "num_warps": [4, 8, 16],
    # "instruction_sched_variant": ["none", "llvm-iglp-1", "llvm-iglp-8", "local-prefetch"],
    "matrix_instr_nonkdim": [16],
    "num_stages": [1, 2, 3],
    "sched_latency": ["none", "mmac5-ds10"],
    "kpack": [1, 2],
})
fn = triton.utils.hcutune(configs=configs, key=key, perf_debug=True)(_attn_fwd)

def attn_false(q, k, v, block_m, block_n, q_scale, k_scale, tensor_layout="HND", attn_mask=None, output_dtype=torch.float16, return_lse=False):
    o = torch.empty(q.shape, dtype=output_dtype, device=q.device)

    if tensor_layout == "HND":
        b, h_qo, qo_len, head_dim = q.shape
        _, h_kv, kv_len, _ = k.shape

        stride_bz_q, stride_h_q, stride_seq_q = q.stride(0), q.stride(1), q.stride(2)
        stride_bz_k, stride_h_k, stride_seq_k = k.stride(0), k.stride(1), k.stride(2)
        stride_bz_v, stride_h_v, stride_seq_v = v.stride(0), v.stride(1), v.stride(2)
        stride_bz_o, stride_h_o, stride_seq_o = o.stride(0), o.stride(1), o.stride(2)
    elif tensor_layout == "NHD":
        b, qo_len, h_qo, head_dim = q.shape
        _, kv_len, h_kv, _ = k.shape

        stride_bz_q, stride_h_q, stride_seq_q = q.stride(0), q.stride(2), q.stride(1)
        stride_bz_k, stride_h_k, stride_seq_k = k.stride(0), k.stride(2), k.stride(1)
        stride_bz_v, stride_h_v, stride_seq_v = v.stride(0), v.stride(2), v.stride(1)
        stride_bz_o, stride_h_o, stride_seq_o = o.stride(0), o.stride(2), o.stride(1)
    else:
        raise ValueError(f"tensor_layout {tensor_layout} not supported")

    if attn_mask is not None:
        stride_bz_mask, stride_h_mask, stride_m_mask, stride_n_mask = attn_mask.stride(0), attn_mask.stride(1), attn_mask.stride(2), attn_mask.stride(3)
    else:
        stride_bz_mask, stride_h_mask, stride_m_mask, stride_n_mask = 0, 0, 0, 0

    HEAD_DIM_K = head_dim
    num_kv_groups = h_qo // h_kv

    if return_lse:
        lse = torch.empty([b, h_qo, qo_len], dtype=torch.float32, device=q.device)
    else:
        lse = torch.empty([0], dtype=torch.float32, device='cpu')

    grid = lambda META: (triton.cdiv(qo_len, block_m), h_qo, b)
    fn[grid](
        q, k, v, q_scale, k_scale, o, attn_mask, lse,
        stride_bz_q, stride_h_q, stride_seq_q, 
        stride_bz_k, stride_h_k, stride_seq_k,  
        stride_bz_v, stride_h_v, stride_seq_v,  
        stride_bz_o, stride_h_o, stride_seq_o,
        stride_bz_mask, stride_h_mask, stride_m_mask, stride_n_mask,
        qo_len, kv_len,
        h_qo, num_kv_groups,
        HEAD_DIM=HEAD_DIM_K,
        RETURN_LSE=return_lse,
        BLOCK_M=block_m,
        BLOCK_N=block_n,
    )

    return o, lse


key = [
    'L',
    'C',
    'BLK',
]
configs = get_triton_configs({
    "waves_per_eu": [1, 2, 4, 8, 16],
    "num_warps": [1, 2, 4, 8, 16],
    "num_stages": [1, 2, 3],
})
fn_quant = triton.utils.hcutune(configs=configs, key=key, perf_debug=True)(quant_per_block_int8_kernel)

def per_block_int8_triton(q, k, km=None, BLKQ=128, BLKK=64, sm_scale=None, tensor_layout="HND"):
    q_int8 = torch.empty(q.shape, dtype=torch.int8, device=q.device)
    k_int8 = torch.empty(k.shape, dtype=torch.int8, device=k.device)

    if km is not None:
        k = k - km

    if tensor_layout == "HND":
        b, h_qo, qo_len, head_dim = q.shape
        _, h_kv, kv_len, _ = k.shape

        stride_bz_q, stride_h_q, stride_seq_q = q.stride(0), q.stride(1), q.stride(2)
        stride_bz_qo, stride_h_qo, stride_seq_qo = q_int8.stride(0), q_int8.stride(1), q_int8.stride(2)
        stride_bz_k, stride_h_k, stride_seq_k = k.stride(0), k.stride(1), k.stride(2)
        stride_bz_ko, stride_h_ko, stride_seq_ko = k_int8.stride(0), k_int8.stride(1), k_int8.stride(2)
    elif tensor_layout == "NHD":
        b, qo_len, h_qo, head_dim = q.shape
        _, kv_len, h_kv, _ = k.shape

        stride_bz_q, stride_h_q, stride_seq_q = q.stride(0), q.stride(2), q.stride(1)
        stride_bz_qo, stride_h_qo, stride_seq_qo = q_int8.stride(0), q_int8.stride(2), q_int8.stride(1)
        stride_bz_k, stride_h_k, stride_seq_k = k.stride(0), k.stride(2), k.stride(1)
        stride_bz_ko, stride_h_ko, stride_seq_ko = k_int8.stride(0), k_int8.stride(2), k_int8.stride(1)
    else:
        raise ValueError(f"Unknown tensor layout: {tensor_layout}")

    q_scale = torch.empty((b, h_qo, (qo_len + BLKQ - 1) // BLKQ), device=q.device, dtype=torch.float32)
    k_scale = torch.empty((b, h_kv, (kv_len + BLKK - 1) // BLKK), device=q.device, dtype=torch.float32)

    if sm_scale is None:
        sm_scale = head_dim**-0.5

    grid = lambda META: ((qo_len + META['BLK'] - 1) // META['BLK'], h_qo, b)
    fn_quant[grid](
        q, q_int8, q_scale, qo_len,
        stride_bz_q, stride_h_q, stride_seq_q,
        stride_bz_qo, stride_h_qo, stride_seq_qo,
        q_scale.stride(0), q_scale.stride(1),
        sm_scale=(sm_scale * 1.44269504),
        C=head_dim, BLK=BLKQ,
    )

    grid = lambda META: ((kv_len + META['BLK'] - 1) // META['BLK'], h_kv, b)
    fn_quant[grid](
        k, k_int8, k_scale, kv_len,
        stride_bz_k, stride_h_k, stride_seq_k,
        stride_bz_ko, stride_h_ko, stride_seq_ko,
        k_scale.stride(0), k_scale.stride(1),
        sm_scale=1.0,
        C=head_dim, BLK=BLKK,
    )

    return q_int8, q_scale, k_int8, k_scale


@triton.utils.graphtune
def sageattn_qk_int8_pv_fp16_triton(
    q: torch.Tensor, 
    k: torch.Tensor, 
    v: torch.Tensor, 
    block_m,
    block_n,
    tensor_layout: str = "HND",
    quantization_backend: str = "triton",
    is_causal: bool =False, 
    attn_mask: Optional[torch.Tensor] = None,
    sm_scale: Optional[float] = None, 
    smooth_k: bool = True,
    return_lse: bool = False,
    **kwargs: Any,
) -> torch.Tensor:
    """
    SageAttention with per-block INT8 quantization for Q and K, FP16 PV with FP16 accumulation, implemented using Triton.
    The FP16 accumulator is added to a FP32 buffer immediately after each iteration.

    Parameters
    ----------
    q : torch.Tensor
        The query tensor. Shape:
        - If `tensor_layout` is "HND": ``[batch_size, num_qo_heads, qo_len, head_dim]``.
        - If `tensor_layout` is "NHD": ``[batch_size, qo_len, num_qo_heads, head_dim]``.

    k : torch.Tensor
        The key tensor. Shape:
        - If `tensor_layout` is "HND": ``[batch_size, num_kv_heads, kv_len, head_dim]``.
        - If `tensor_layout` is "NHD": ``[batch_size, kv_len, num_kv_heads, head_dim]``.

    v : torch.Tensor
        The value tensor. Shape:
        - If `tensor_layout` is "HND": ``[batch_size, num_kv_heads, kv_len, head_dim]``.
        - If `tensor_layout` is "NHD": ``[batch_size, kv_len, num_kv_heads, head_dim]``.

    tensor_layout : str
        The tensor layout, either "HND" or "NHD".
        Default: "HND".

    quantization_backend : str
        The quantization backend, either "triton" or "cuda".
        "cuda" backend offers better performance due to kernel fusion.

    is_causal : bool
        Whether to apply causal mask to the attention matrix. Only applicable when qo_len == kv_len.
        Default: False.

    attn_mask : Optional[torch.Tensor]
        The attention mask tensor, of dtype bool or float32.
        Should be able to broadcast to the shape of the matrix qk^T.
        Default: None.

    sm_scale : Optional[float]
        The scale used in softmax, if not provided, will be set to ``1.0 / sqrt(head_dim)``.

    smooth_k : bool
        Whether to smooth the key tensor by subtracting the mean along the sequence dimension.
        Default: True.

    return_lse : bool
        Whether to return the log sum of the exponentiated attention weights. Used for cases like Ring Attention.
        Default: False.

    Returns
    -------
    torch.Tensor
        The output tensor. Shape:
        - If `tensor_layout` is "HND": ``[batch_size, num_qo_heads, qo_len, head_dim]``.
        - If `tensor_layout` is "NHD": ``[batch_size, qo_len, num_qo_heads, head_dim]``.

    torch.Tensor
        The logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax normalization factor).
        Shape: ``[batch_size, num_qo_heads, qo_len]``.
        Only returned if `return_lse` is True.

    Note
    ----
    - ``num_qo_heads`` must be divisible by ``num_kv_heads``. 
    - The tensors `q`, `k`, and `v` must have the dtype ``torch.float16``, ``torch.bfloat16`` or ``torch.float32``.
    - All tensors must be on the same cuda device.
    - `smooth_k` will introduce slight overhead but will improve the accuracy under most circumstances.
    """

    dtype = q.dtype
    # assert q.is_cuda, "Input tensors must be on cuda."
    assert dtype in [torch.float16, torch.bfloat16], "Input tensors must be in dtype of torch.float16 or torch.bfloat16"
    assert q.device == k.device == v.device, "All tensors must be on the same device."
    assert q.dtype == k.dtype == v.dtype, "All tensors must have the same dtype."

    if attn_mask is not None:
        assert attn_mask.dtype == torch.bool or attn_mask.dtype == q.dtype, "attn_mask must be of dtype bool or the same dtype as q."
        assert attn_mask.device == q.device, "All tensors must be on the same device."

    # FIXME(DefTruth): make sage attention work compatible with distributed 
    # env, for example, xDiT which launch by torchrun. Without this workaround, 
    # sage attention will run into illegal memory access error after first 
    # inference step in distributed env for multi gpus inference. This small
    # workaround also make sage attention work compatible with torch.compile
    # through non-fullgraph compile mode.
    # torch.cuda.set_device(v.device)

    head_dim_og = q.size(-1)

    if head_dim_og < 64:
        q = torch.nn.functional.pad(q, (0, 64 - head_dim_og))
        k = torch.nn.functional.pad(k, (0, 64 - head_dim_og))
        v = torch.nn.functional.pad(v, (0, 64 - head_dim_og))
    elif head_dim_og > 64 and head_dim_og < 128:
        q = torch.nn.functional.pad(q, (0, 128 - head_dim_og))
        k = torch.nn.functional.pad(k, (0, 128 - head_dim_og))
        v = torch.nn.functional.pad(v, (0, 128 - head_dim_og))
    elif head_dim_og > 128:
        raise ValueError(f"Unsupported head_dim: {head_dim_og}")

    # assert last dim is contiguous
    assert q.stride(-1) == 1 and k.stride(-1) == 1 and v.stride(-1) == 1, "Last dim of qkv must be contiguous."

    seq_dim = 1 if tensor_layout == "NHD" else 2
    nh_dim = 2 if tensor_layout == 0 else 1

    if smooth_k:
        km = k.mean(dim=seq_dim, keepdim=True)
        k = k - km
        nqheads = q.size(2)
        nkheads = k.size(2)
        q_per_kv_heads = nqheads // nkheads
        if q_per_kv_heads > 1:
            # nheads_k => nheads_q
            km_broadcast = torch.repeat_interleave(km, q_per_kv_heads, dim=nh_dim)
        else:
            km_broadcast = km
        if return_lse:
            if tensor_layout == "NHD":
                lse_correction = torch.matmul(q.transpose(1, 2), km_broadcast.transpose(1, 2).transpose(2, 3)).squeeze(-1).to(torch.float32)
            else:
                lse_correction = torch.matmul(q, km_broadcast.transpose(2, 3)).squeeze(-1).to(torch.float32)
    else:
        km = None

    if dtype == torch.bfloat16 or dtype == torch.float32:
        v = v.to(torch.float16)

    if sm_scale is None:
        sm_scale = 1.0 / (head_dim_og ** 0.5)

    if quantization_backend == "triton":
        q_int8, q_scale, k_int8, k_scale = per_block_int8_triton(q, k, km=km, BLKQ=block_m, BLKK=block_n, sm_scale=sm_scale, tensor_layout=tensor_layout)
    else:
        raise ValueError(f"Unsupported quantization backend: {quantization_backend}")

    if is_causal:
        assert attn_mask is None, "Mask should be None for causal attention."
        o, lse = attn_true(q_int8, k_int8, v, block_m, block_n, q_scale, k_scale, tensor_layout=tensor_layout, output_dtype=dtype, return_lse=return_lse)
    else:
        if attn_mask is not None:
            if tensor_layout == "HND":
                target_shape = (q.shape[0], q.shape[1], q.shape[2], k.shape[2])
            elif tensor_layout == "NHD":
                target_shape = (q.shape[0], q.shape[2], q.shape[1], k.shape[1])
            else:
                raise ValueError(f"tensor_layout {tensor_layout} not supported")
            try:
                attn_mask = attn_mask.expand(target_shape)
            except Exception:
                raise AssertionError(f"attn_mask shape {attn_mask.shape} cannot be broadcast to {target_shape}")
        o, lse = attn_false(q, k, v, block_m, block_n, q_scale, k_scale, tensor_layout=tensor_layout, output_dtype=dtype,
                            attn_mask=attn_mask, return_lse=return_lse)

    o = o[..., :head_dim_og]

    if return_lse:
        return o, lse / 1.44269504 + lse_correction * sm_scale if smooth_k else lse / 1.44269504
    else:
        return o


PROBLEM_SIZES = [
    # "batch, num_qo_heads, qo_len, num_kv_heads, kv_len, head_dim, tensor_layout, is_causal"
    (1, 14040, 40, 14040, 40, 128, "NHD", False),
    (1, 14040, 40, 512, 40, 128, "NHD", False),
    (2, 14040, 40, 14040, 40, 128, "NHD", False),
    (2, 14040, 40, 512, 40, 128, "NHD", False),
    (2, 24, 1280, 24, 1280, 128, "HND", False),
    (1, 24, 1280, 24, 1280, 128, "HND", False),
    (1, 24, 4352, 24, 4352, 128, "HND", False),
    (2, 24, 4352, 24, 4352, 128, "HND", False),
    (1, 19440, 1, 19440, 1, 128, "NHD", False),
    (1, 2430, 40, 512, 40, 128, "NHD", False),
    (1, 57600, 1, 57600, 1, 128, "NHD", False),
    (1, 7200, 40, 512, 40, 128, "NHD", False),
    (1, 14080, 1, 14080, 1, 128, "NHD", False),
    (1, 3520, 24, 512, 24, 128, "NHD", False),
]

X_VALS = [(*sizes, block_m, block_n) for sizes in PROBLEM_SIZES \
          for block_m in [32, 64, 128, 256] \
          for block_n in [32, 64, 128, 256] \
          if block_m != 256 and block_n != 256]

configs = [
    triton.testing.Benchmark(
        x_names=['batch', 'num_qo_heads', 'qo_len', 'num_kv_heads', 'kv_len', 'head_dim', 'tensor_layout', 'is_causal', 'block_m', 'block_n'],
        x_vals=X_VALS,
        line_arg='provider',
        line_vals=['triton'],
        line_names=['Triton'],
        styles=[('red', '-')],
        ylabel='Latency',
        xlabel='sizes',
        plot_name='qk_int8_pv_fp16_performance',
        args={'dtype': torch.float16},
    )
]


@triton.utils.dist_perf_report(configs)
def bench_qk_int8_pv_fp16(batch, num_qo_heads, qo_len, num_kv_heads, kv_len, head_dim,
                          tensor_layout, is_causal, block_m, block_n, provider, dtype=torch.float16):
    print(
        f"Tuning batch: {batch}, num_qo_heads: {num_qo_heads}, qo_len: {qo_len}, "
        f"num_kv_heads: {num_kv_heads}, kv_len: {kv_len}, head_dim: {head_dim}, "
        f"tensor_layout: {tensor_layout}, BLOCK_M: {block_m}, BLOCK_N: {block_n}"
    )

    device = "cuda" if not os.getenv("TRITON_HCUTUNE_COMPILE_ONLY", "0") == "1" else "cpu"

    q = torch.randn((batch, num_qo_heads, qo_len, head_dim), dtype=dtype, device=device)
    k = torch.randn((batch, num_kv_heads, kv_len, head_dim), dtype=dtype, device=device)
    v = torch.randn((batch, num_kv_heads, kv_len, head_dim), dtype=dtype, device=device)

    fn = lambda: sageattn_qk_int8_pv_fp16_triton(q, k, v, block_m, block_n, tensor_layout=tensor_layout,
                                                 is_causal=is_causal)
    return triton.testing.do_bench(fn)


if __name__ == "__main__":
    print(f"Triton QK Int8 PV FP16 Fine-tuning")

    bench_qk_int8_pv_fp16.run(print_data=True, save_path="bench_qk_int8_pv_fp16_out")

