import torch
import torch.nn.functional as F
import triton
import triton.language as tl
       
@triton.autotune(
    configs=[
        triton.Config({"BLOCK_N": BLOCK_N}, num_warps=num_warps, num_stages=1)
        for num_warps in [1, 2, 4, 8]
        for BLOCK_N in [32, 64]
    ],
    key=['local_block_num', 'head_dim'],
)
@triton.jit
def block_attn_decoding_kernel(
    q,                             # pointer to float32, shape [batch_size, n_head * gqa_size, head_dim]
    k_min,                         # pointer to float32, shape [batch_size, num_blocks, n_head, head_dim]
    k_max,                         # pointer to float32, shape [batch_size, num_blocks, n_head, head_dim]
    num_blocks,                    # pointer to int32, shape [batch_size]
    local_block_num: tl.constexpr, # int32, number of local blocks per query
    head_dim: tl.constexpr,        # int32, dimension per head
    stride_qb, stride_qh, stride_qd,
    stride_kb, stride_kt, stride_kh, stride_kd,
    stride_ab, stride_ah, stride_at,
    attn_score,                    # pointer to int32, shape [batch_size, n_head, num_blocks]
    BLOCK_N: tl.constexpr,         # int32, number of queries per block
):
    batch_id = tl.program_id(0)
    head_idx = tl.program_id(1)
    start_block_idx = tl.program_id(2) * BLOCK_N
    base_q_ptr = tl.make_block_ptr(q + batch_id * stride_qb + head_idx * stride_qh, shape=(head_dim,), block_shape=(head_dim,), strides=(stride_qd,), offsets=(0,), order=(0,))
    mean_q = tl.load(base_q_ptr).to(tl.float32)
    POS_INF = 1e38

    total_n_blocks = tl.load(num_blocks + batch_id)
    block_block_idx = start_block_idx + tl.arange(0, BLOCK_N)
    k_min_ptr = tl.make_block_ptr(k_min + batch_id * stride_kb + head_idx * stride_kh, shape=(total_n_blocks, head_dim), block_shape=(BLOCK_N, head_dim), strides=(stride_kt, stride_kd), offsets=(start_block_idx, 0), order=(0, 1))
    k_max_ptr = tl.make_block_ptr(k_max + batch_id * stride_kb + head_idx * stride_kh, shape=(total_n_blocks, head_dim), block_shape=(BLOCK_N, head_dim), strides=(stride_kt, stride_kd), offsets=(start_block_idx, 0), order=(0, 1))
    k_min_val = tl.load(k_min_ptr, boundary_check=(0, 1)).to(tl.float32)
    k_max_val = tl.load(k_max_ptr, boundary_check=(0, 1)).to(tl.float32)
    score = tl.maximum(k_min_val * mean_q, k_max_val * mean_q)
    score = tl.sum(score, axis=1)
    score = tl.where(block_block_idx >= total_n_blocks - local_block_num, POS_INF, score)
    attn_score_ptr = tl.make_block_ptr(attn_score + batch_id * stride_ab + head_idx * stride_ah, shape=(total_n_blocks,), block_shape=(BLOCK_N,), strides=(stride_at,), offsets=(start_block_idx,), order=(0,))
    tl.store(attn_score_ptr, score, boundary_check=(0,))

def block_attn_decoding(q, k_min, k_max, num_blocks, local_block_num):
    assert num_blocks.max().item() <= k_min.shape[1] , "num_blocks should be less than or equal to k_min.shape[1]"
    min_val = torch.finfo(torch.float32).min
    batch, n_kv_heads, head_dim = q.shape[0], k_min.shape[2], k_min.shape[3]
    attn_score = torch.full((batch, n_kv_heads, num_blocks.max().item()), fill_value=min_val, device=q.device, dtype=torch.float32)
    grid = lambda META: (batch, n_kv_heads, triton.cdiv(num_blocks.max().item(), META['BLOCK_N']))
    with torch.cuda.device(q.device.index): 
        block_attn_decoding_kernel[grid](q, k_min, k_max, 
            num_blocks, local_block_num, head_dim,
            *q.stride(),
            *k_min.stride(),
            *attn_score.stride(),
            attn_score,
        )
    return attn_score

@torch.compile(fullgraph=True)
def get_topk_indices(block_attn_score, block_index_mask, max_num_selected_blocks):
    block_index = torch.topk(block_attn_score, k=max_num_selected_blocks, dim=-1, sorted=True).indices
    topk_indices = block_index.masked_fill_(~block_index_mask, -1)
    topk_indices = torch.sort(topk_indices, dim=-1, descending=False).values
    return topk_indices.to(torch.int32)

@torch.compile(fullgraph=True)
def get_num_blocks(cache_seqlens, block_size, sparse_ratio, min_block_num):
    num_blocks = (cache_seqlens + (block_size - 1)) // block_size
    num_selected_blocks = torch.ceil(num_blocks * sparse_ratio)
    num_selected_blocks = torch.maximum(num_selected_blocks, num_blocks.clamp(max=min_block_num)).long()
    return num_blocks, num_selected_blocks

class KVManager:
    def __init__(self, num_heads, block_size, sparse_ratio, local_block_num, min_block_num):
        self.num_heads = num_heads
        self.block_size = block_size
        self.sparse_ratio = sparse_ratio
        self.local_block_num = local_block_num
        self.min_block_num = min_block_num
        self.block_max_key = None
        self.block_min_key = None
        self.num_elements = None
        assert self.local_block_num <= self.min_block_num, "local_block_num should be less than or equal to min_block_num"

    def init_centeroids(self, key, cache_seqlens):
        bsz, seqlen, num_heads, head_dim = key.shape
        max_val = torch.finfo(key.dtype).max
        min_val = torch.finfo(key.dtype).min
        self.num_elements = seqlen
        key_block = key.reshape(bsz, -1, self.block_size, num_heads, head_dim)
        self.block_max_key = key_block.max(dim=2).values
        self.block_min_key = key_block.min(dim=2).values
        num_blocks = (cache_seqlens + self.block_size - 1) // self.block_size
        key_last_block = key_block[torch.arange(bsz, device=key.device), num_blocks - 1]
        valid_mask = torch.arange(self.block_size, device=key.device) + (num_blocks[:, None] - 1) * self.block_size < cache_seqlens[:, None]
        last_min_key = torch.masked_fill(key_last_block, ~valid_mask[:, :, None, None], max_val).min(dim=1).values
        last_max_key = torch.masked_fill(key_last_block, ~valid_mask[:, :, None, None], min_val).max(dim=1).values
        self.block_max_key[torch.arange(bsz, device=key.device), num_blocks - 1] = last_max_key
        self.block_min_key[torch.arange(bsz, device=key.device), num_blocks - 1] = last_min_key
        causal_mask = torch.arange(self.block_max_key.shape[1], device=key.device) < num_blocks[:, None]
        self.block_max_key.masked_fill_(~causal_mask[:, :, None, None], min_val)
        self.block_min_key.masked_fill_(~causal_mask[:, :, None, None], max_val)
        
    @torch.compile(fullgraph=True)
    def update_centeroids(self, key, cache_seqlens):
        num_blocks = (cache_seqlens + self.block_size - 1) // self.block_size
        batch_index = torch.arange(num_blocks.shape[0], device=key.device)
        self.block_max_key[batch_index, num_blocks - 1] = torch.maximum(self.block_max_key[batch_index, num_blocks - 1], key)
        self.block_min_key[batch_index, num_blocks - 1] = torch.minimum(self.block_min_key[batch_index, num_blocks - 1], key)

    def clear_centeroids(self):
        self.block_max_key = None
        self.block_min_key = None
        self.num_elements = None

    def get_kv_cache_indices(self, query, cache_seqlens):
        bsz, num_heads, head_dim = query.shape
        max_val = torch.finfo(torch.float32).max
        query = query.view(bsz, 1, self.num_heads, num_heads // self.num_heads, head_dim).mean(dim=3).float() * (head_dim ** -0.5)
        attn_score = torch.maximum(query * self.block_max_key.float(), query * self.block_min_key.float()).sum(dim=-1)
        topk_indices = torch.full((bsz, self.num_heads, self.block_max_key.shape[1]), device=query.device, dtype=torch.int32, fill_value=-1)
        num_blocks = (cache_seqlens + self.block_size - 1) // self.block_size
        num_selected_blocks = torch.ceil(num_blocks * self.sparse_ratio)[:, None].repeat(1, self.num_heads)

        num_selected_blocks = torch.maximum(num_selected_blocks, num_blocks.clamp(max=self.min_block_num)[:, None]).long()
        for b in range(bsz):
            local_attn_score = attn_score[b, :num_blocks[b]]
            local_attn_score[-self.local_block_num:] = max_val
            block_index = torch.sort(local_attn_score, dim=0, descending=True).indices.transpose(0, 1)
            block_index_mask = torch.arange(num_blocks[b], device=query.device) < num_selected_blocks[b][:, None]
            topk_indices[b, :, :num_blocks[b]] = block_index.masked_fill(~block_index_mask, -1).to(torch.int32)
        topk_indices = torch.sort(topk_indices, dim=-1, descending=True).values
        return topk_indices

    def get_kv_cache_indices_fast(self, query, cache_seqlens):
        bsz, num_heads, head_dim = query.shape
        num_blocks, num_selected_blocks = get_num_blocks(cache_seqlens, self.block_size, self.sparse_ratio, self.min_block_num)
        max_num_selected_blocks = num_selected_blocks.max().item()
        query = query.view(bsz, self.num_heads, num_heads // self.num_heads, head_dim).mean(dim=2) * (head_dim ** -0.5)
        block_attn_score = block_attn_decoding(query, self.block_max_key, self.block_min_key, num_blocks, self.local_block_num)
        block_index_mask = torch.arange(max_num_selected_blocks, device=query.device) < num_selected_blocks[:, None, None]
        topk_indices = get_topk_indices(block_attn_score, block_index_mask, max_num_selected_blocks)
        return topk_indices