import torch

from kernels.benchmark import Benchmark


def ref_masked_attention(
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    scale: float,
) -> torch.Tensor:
    # query: (q, h, d), key: (k, h, d), value: (k, h, d)
    # Transpose to (h, q, d) and (h, k, d) for batched matmul
    q = query.transpose(0, 1)  # (h, q, d)
    k = key.transpose(0, 1)  # (h, k, d)
    v = value.transpose(0, 1)  # (h, k, d)

    # Compute attention scores: (h, q, d) @ (h, d, k) -> (h, q, k)
    attn_weights = (scale * torch.matmul(q, k.transpose(-1, -2))).float()
    attn_weights = torch.softmax(attn_weights, dim=-1).to(value.dtype)

    # Compute output: (h, q, k) @ (h, k, d) -> (h, q, d)
    out = torch.matmul(attn_weights, v)

    # Transpose back to (q, h, d)
    return out.transpose(0, 1)


def ref_paged_attention(
    query: torch.Tensor,
    key_cache: torch.Tensor,
    value_cache: torch.Tensor,
    block_tables: torch.Tensor,
    seq_lens: torch.Tensor,
    scale: float,
) -> torch.Tensor:
    num_seqs = query.shape[0]
    num_heads = query.shape[1]
    head_size = query.shape[2]
    block_size = value_cache.shape[3]
    max_seq_len = int(seq_lens.max().item())

    # Create position indices for all sequences up to max_seq_len
    positions = torch.arange(max_seq_len, device=query.device)
    block_indices = positions // block_size  # (max_seq_len,)
    block_offsets = positions % block_size  # (max_seq_len,)

    # Gather block numbers for all sequences: (num_seqs, max_seq_len)
    block_numbers = block_tables[:, block_indices.long()]

    # Flatten for gathering: (num_seqs * max_seq_len,)
    flat_block_numbers = block_numbers.reshape(-1)
    flat_offsets = block_offsets.repeat(num_seqs)

    # Gather keys: key_cache is (num_blocks, num_heads, head_size // x, block_size, x)
    # Index into [block_number, :, :, offset, :] and reshape
    keys = key_cache[flat_block_numbers, :, :, flat_offsets, :]
    keys = keys.reshape(num_seqs, max_seq_len, num_heads, head_size)
    keys = keys.transpose(1, 2)  # (num_seqs, num_heads, max_seq_len, head_size)

    # Gather values: value_cache is (num_blocks, num_heads, head_size, block_size)
    values = value_cache[flat_block_numbers, :, :, flat_offsets]
    values = values.reshape(num_seqs, max_seq_len, num_heads, head_size)
    values = values.transpose(1, 2)  # (num_seqs, num_heads, max_seq_len, head_size)

    # Query: (num_seqs, num_heads, head_size) -> (num_seqs, num_heads, 1, head_size)
    q = query.unsqueeze(2)

    # Compute attention scores: (num_seqs, num_heads, 1, head_size) @ (num_seqs, num_heads, head_size, max_seq_len)
    attn_weights = (scale * torch.matmul(q, keys.transpose(-1, -2))).float()

    # Create causal mask for variable sequence lengths
    # Mask out positions beyond seq_len for each sequence
    seq_mask = positions.unsqueeze(0) >= seq_lens.unsqueeze(
        1
    )  # (num_seqs, max_seq_len)
    seq_mask = seq_mask.unsqueeze(1).unsqueeze(2)  # (num_seqs, 1, 1, max_seq_len)
    attn_weights = attn_weights.masked_fill(seq_mask, float("-inf"))

    attn_weights = torch.softmax(attn_weights, dim=-1).to(values.dtype)

    # Compute output: (num_seqs, num_heads, 1, max_seq_len) @ (num_seqs, num_heads, max_seq_len, head_size)
    out = torch.matmul(attn_weights, values)

    return out.squeeze(2)  # (num_seqs, num_heads, head_size)


class PagedAttentionBenchmark(Benchmark):
    seed: int = 42

    def setup(self):
        num_seqs = 4
        num_heads = 8
        head_size = 64
        block_size = 16
        max_seq_len = 128
        num_blocks = 64
        dtype = torch.float16

        self.num_heads = num_heads
        self.block_size = block_size
        self.max_seq_len = max_seq_len
        self.scale = 1.0 / (head_size**0.5)

        # Query tensor (current token)
        self.query = torch.randn(
            num_seqs, num_heads, head_size, device=self.device, dtype=dtype
        )

        # KV cache with proper layout for the kernel
        # x = 16 // element_size, for float16 x = 8
        x = 16 // torch.tensor([], dtype=dtype).element_size()
        self.key_cache = torch.randn(
            num_blocks,
            num_heads,
            head_size // x,
            block_size,
            x,
            device=self.device,
            dtype=dtype,
        )
        self.value_cache = torch.randn(
            num_blocks,
            num_heads,
            head_size,
            block_size,
            device=self.device,
            dtype=dtype,
        )

        # Block tables: mapping from sequences to memory blocks
        max_num_blocks_per_seq = (max_seq_len + block_size - 1) // block_size
        self.block_tables = torch.randint(
            0,
            num_blocks,
            (num_seqs, max_num_blocks_per_seq),
            device=self.device,
            dtype=torch.int32,
        )

        # Sequence lengths
        self.seq_lens = torch.tensor(
            [64, 96, 48, 128], device=self.device, dtype=torch.int32
        )

        # KV scales
        self.k_scale = torch.tensor(1.0, dtype=torch.float32, device=self.device)
        self.v_scale = torch.tensor(1.0, dtype=torch.float32, device=self.device)

        # Output tensor
        self.out = torch.empty_like(self.query)

    def benchmark_base(self):
        self.kernel.paged_attention_v1(
            self.out,
            self.query,
            self.key_cache,
            self.value_cache,
            num_kv_heads=self.num_heads,
            scale=self.scale,
            block_tables=self.block_tables,
            seq_lens=self.seq_lens,
            block_size=self.block_size,
            max_seq_len=self.max_seq_len,
            alibi_slopes=None,
            kv_cache_dtype="auto",
            k_scale=self.k_scale,
            v_scale=self.v_scale,
        )

    def verify_base(self) -> torch.Tensor:
        return ref_paged_attention(
            self.query,
            self.key_cache,
            self.value_cache,
            self.block_tables,
            self.seq_lens,
            self.scale,
        )

    def setup_large(self):
        num_seqs = 16
        num_heads = 32
        head_size = 128
        block_size = 16
        max_seq_len = 512
        num_blocks = 256
        dtype = torch.float16

        self.num_heads = num_heads
        self.block_size = block_size
        self.max_seq_len = max_seq_len
        self.scale = 1.0 / (head_size**0.5)

        self.query = torch.randn(
            num_seqs, num_heads, head_size, device=self.device, dtype=dtype
        )

        x = 16 // torch.tensor([], dtype=dtype).element_size()
        self.key_cache = torch.randn(
            num_blocks,
            num_heads,
            head_size // x,
            block_size,
            x,
            device=self.device,
            dtype=dtype,
        )
        self.value_cache = torch.randn(
            num_blocks,
            num_heads,
            head_size,
            block_size,
            device=self.device,
            dtype=dtype,
        )

        max_num_blocks_per_seq = (max_seq_len + block_size - 1) // block_size
        self.block_tables = torch.randint(
            0,
            num_blocks,
            (num_seqs, max_num_blocks_per_seq),
            device=self.device,
            dtype=torch.int32,
        )

        # Variable sequence lengths
        self.seq_lens = torch.randint(
            64, max_seq_len + 1, (num_seqs,), device=self.device, dtype=torch.int32
        )

        self.k_scale = torch.tensor(1.0, dtype=torch.float32, device=self.device)
        self.v_scale = torch.tensor(1.0, dtype=torch.float32, device=self.device)

        self.out = torch.empty_like(self.query)

    def benchmark_large(self):
        self.kernel.paged_attention_v1(
            self.out,
            self.query,
            self.key_cache,
            self.value_cache,
            num_kv_heads=self.num_heads,
            scale=self.scale,
            block_tables=self.block_tables,
            seq_lens=self.seq_lens,
            block_size=self.block_size,
            max_seq_len=self.max_seq_len,
            alibi_slopes=None,
            kv_cache_dtype="auto",
            k_scale=self.k_scale,
            v_scale=self.v_scale,
        )

    def verify_large(self) -> torch.Tensor:
        return ref_paged_attention(
            self.query,
            self.key_cache,
            self.value_cache,
            self.block_tables,
            self.seq_lens,
            self.scale,
        )