import math

import pytest
import torch

import flag_gems

from .performance_utils import GenericBenchmark, SkipVersion


def torch_flash_fwd_cpu(
    q,
    k,
    v,
    scale,
    is_causal=False,
    dropout_p=0,
    return_debug_mask=False,
    window_size_left=None,
    window_size_right=None,
    **unused_kwargs,
):
    # Use the basic PyTorch operators (matmul + masked_fill + softmax + logsumexp) to implement attention,
    # **avoid the BLAS path in SDPA**, and **one matmul** is more stable.
    # Input layout (B, H, S_q, D). The q_shape generated by make_input is (batch, num_head, q_seq_len, head_size).
    B, H, S_q, D = q.shape
    _, _, S_k, _ = k.shape
    W_L = window_size_left if window_size_left is not None else -1
    W_R = window_size_right if window_size_right is not None else -1
    if is_causal:
        W_R = 0
        if W_L < 0 or W_L > S_q:
            W_L = S_q

    dtype_og = q.dtype

    # Construct the SWA + causal attn_mask (S_q, S_k).
    # mask[i, j] = True indicates that the value is **reserved** (not masked),
    # and False indicates that the value is masked.
    q_idx = torch.arange(S_q, device=q.device).view(S_q, 1)
    k_idx = torch.arange(S_k, device=q.device).view(1, S_k)
    sk_minus_sq = S_k - S_q
    keep = torch.ones(S_q, S_k, dtype=torch.bool, device=q.device)
    if is_causal:
        # causal: k_idx <= q_idx + sk_minus_sq + W_R (note that W_R = 0 because is_causal)
        keep = keep & (k_idx <= q_idx + sk_minus_sq + W_R)
    else:
        if W_L >= 0:
            # W_L indicates the size of the left window,
            # q[i] focuses on k[i + sk_minus_sq - W_L : i + sk_minus_sq + 1]
            keep = keep & (k_idx >= q_idx + sk_minus_sq - W_L)
        if W_R >= 0:
            keep = keep & (k_idx <= q_idx + sk_minus_sq + W_R)
    # Broadcast to (1, 1, S_q, S_k), aligned with scores.
    attn_mask = keep.view(1, 1, S_q, S_k)  # (1, 1, S_q, S_k)

    # Compute scores (B, H, S_q, S_k) in fp32 at a time
    # Use fp32 to compute scores to improve the precision.
    # (Compared with the fp16 accumulation path of the kernel, the ref has higher precision.)
    scores = torch.matmul(q.float(), k.float().transpose(-2, -1)) * scale  # fp32
    scores = scores.masked_fill(~attn_mask, float("-inf"))  # 应用 mask

    # Softmax along the S_k dimension.
    probs = torch.softmax(scores, dim=-1)  # fp32 (B, H, S_q, S_k)
    # Handle NaN values in fully masked rows (softmax 0/0 = NaN) and replace them with 0.
    probs = torch.nan_to_num(probs, nan=0.0, posinf=0.0, neginf=0.0)

    if dropout_p > 0:
        # Inverted dropout, matching the kernel's semantics: keep each attention
        # weight with probability (1-dropout_p), scale kept entries by 1/(1-dropout_p),
        # and drop the rest to 0 (so the expected sum is preserved). The debug
        # softmax encodes dropped entries as negative (sign-bit convention used by
        # the kernel), so callers can count the drop ratio via (debug < 0)/(debug != 0).
        keep_mask = torch.rand(probs.shape, device=probs.device) < (1.0 - dropout_p)
        P = torch.where(
            keep_mask, probs / (1.0 - dropout_p), torch.zeros_like(probs)
        )
        debug_softmax = torch.where(keep_mask, probs, -probs)
    else:
        P = probs
        debug_softmax = probs

    # Key: Truncate P to the original dtype (consistent with the kernel behavior).
    P = P.to(dtype_og)
    # V is also truncated to the original dtype.
    v_typed = v.to(dtype_og)
    # Whole block P @ V (B, H, S_q, S_k) @ (B, H, S_k, D) -> (B, H, S_q, D)
    out = torch.matmul(P, v_typed)  # fp16/bf16 matmul -> fp32 acc

    # lse = logsumexp(scores, dim=-1) -> (B, H, S_q)
    # The torch.logsumexp function uses the max-subtract technique internally.
    # When all values are -inf, -inf is returned (correct).
    lse = torch.logsumexp(scores, dim=-1)
    lse = torch.nan_to_num(lse, nan=float("-inf"), posinf=float("-inf"), neginf=float("-inf"))

    # The shape of out is (B, H, S_q, D), and it is reshaped back to (B, S_q, H, D) to be consistent with gems_out.
    out = out.transpose(1, 2).to(dtype_og)  # (B, S_q, H, D)

    seed = torch.zeros((), dtype=torch.int64, device=q.device)
    offset = torch.zeros((), dtype=torch.int64, device=q.device)
    if not return_debug_mask:
        # debug_softmax is only meaningful when requested; otherwise return an empty tensor.
        debug_softmax = torch.empty(0, device=q.device, dtype=q.dtype)
    return out, lse, seed, offset, debug_softmax


def torch_flash_attention_forward(
    q, k, v, scale, is_causal, dropout_p=0.0, return_debug_mask=False, **extra_kwargs
):
    if not q.is_cuda:
        # Benchmark tensors use (B, T, H, D) layout (transposed from the kernel's
        # (B, H, T, D)), so transpose back to (B, H, T, D) for the CPU reference.
        q = q.transpose(1, 2)
        k = k.transpose(1, 2)
        v = v.transpose(1, 2)
        # GQA: q may have more heads (H) than k/v (H_k). The CPU reference
        # torch_flash_fwd_cpu expects matching head counts, otherwise the batched
        # matmul q @ k^T fails to broadcast (B,H) against (B,H_k). Expand k/v to
        # the same number of heads as q before calling the reference.
        if q.shape[1] != k.shape[1]:
            g = q.shape[1] // k.shape[1]
            k = k.repeat_interleave(g, dim=1)
            v = v.repeat_interleave(g, dim=1)
        return torch_flash_fwd_cpu(
            q,
            k,
            v,
            scale,
            is_causal=is_causal,
            dropout_p=dropout_p,
            return_debug_mask=return_debug_mask,
            **extra_kwargs,
        )

    return torch.ops.aten._flash_attention_forward(
        q,
        k,
        v,
        None,
        None,
        q.shape[-3],
        k.shape[-3],
        dropout_p,
        is_causal,
        return_debug_mask,
        scale=scale,
        **extra_kwargs,
    )


def gems_flash_attention_forward(
    q, k, v, scale, is_causal, dropout_p=0.0, return_debug_mask=False, **extra_kwargs
):
    return flag_gems.ops.flash_attention_forward(
        q,
        k,
        v,
        None,
        None,
        q.shape[-3],
        k.shape[-3],
        dropout_p,
        is_causal,
        return_debug_mask,
        scale=scale,
        **extra_kwargs,
    )


def torch_flash_attention_supports_alibi(device: str) -> bool:
    if device == "cpu" or not torch.cuda.is_available():
        return False

    try:
        q = torch.randn((1, 16, 1, 64), device=device, dtype=torch.float16)
        k = torch.randn((1, 16, 1, 64), device=device, dtype=torch.float16)
        v = torch.randn((1, 16, 1, 64), device=device, dtype=torch.float16)
        scale = float(1.0 / math.sqrt(64))
        alibi_slopes = torch.ones((1, 1), device=device, dtype=torch.float32) * 0.3
        torch.ops.aten._flash_attention_forward(
            q,
            k,
            v,
            None,
            None,
            q.shape[-3],
            k.shape[-3],
            0.0,
            False,
            False,
            scale=scale,
            alibi_slopes=alibi_slopes,
        )
        return True
    except RuntimeError as e:
        if "does not support alibi" in str(e).lower():
            return False
        raise


class FlashAttentionForwardBenchmark(GenericBenchmark):
    def set_shapes(self, shape_file_path=None):
        self.shapes = []
        for head_size in (64, 128, 192, 256):
            for is_causal in (False, True):
                self.shapes.append(
                    (
                        4,
                        8,
                        8,
                        1024,
                        128,
                        head_size,
                        is_causal,
                        0.0,
                        False,
                        None,
                        None,
                        False,
                    )
                )

        for batch, num_head, q_seq_len, kv_seq_len in (
            (1, 1, 128, 2048),
            (4, 8, 17, 1030),
        ):
            for is_causal in (False, True):
                self.shapes.append(
                    (
                        batch,
                        num_head,
                        num_head,
                        q_seq_len,
                        kv_seq_len,
                        128,
                        is_causal,
                        0.0,
                        False,
                        None,
                        None,
                        False,
                    )
                )

        supports_alibi = torch_flash_attention_supports_alibi(self.device)
        if supports_alibi:
            # GQA + alibi cases
            for head_size in (128, 192):
                for is_causal in (False, True):
                    self.shapes.append(
                        (
                            4,
                            8,
                            2,
                            1024,
                            1024,
                            head_size,
                            is_causal,
                            0.0,
                            False,
                            None,
                            None,
                            True,
                        )
                    )
            for is_causal in (False, True):
                self.shapes.append(
                    (4, 4, 4, 1, 519, 128, is_causal, 0.0, False, None, None, True)
                )

        # Split-KV like cases (q_seq_len=1, num_head_k < num_head).
        for is_causal in (False, True):
            self.shapes.append(
                (1, 4, 1, 1, 1024, 128, is_causal, 0.0, False, None, None, False)
            )
            if supports_alibi:
                self.shapes.append(
                    (1, 4, 1, 1, 1024, 128, is_causal, 0.0, False, None, None, True)
                )

        # Sliding window attention.
        for batch, num_head, q_seq_len, kv_seq_len in (
            (1, 1, 128, 2048),
            (8, 32, 1024, 1024),
            (8, 32, 1024, 128),
            (8, 32, 17, 1030),
        ):
            for window_size_left, window_size_right in ((256, 0), (128, 128)):
                self.shapes.append(
                    (
                        batch,
                        num_head,
                        num_head,
                        q_seq_len,
                        kv_seq_len,
                        128,
                        False,
                        0.0,
                        False,
                        window_size_left,
                        window_size_right,
                        False,
                    )
                )
        self.shapes.append(
            (8, 32, 32, 1024, 1024, 192, False, 0.0, False, 256, 0, False)
        )

        for is_causal in (False, True):
            self.shapes.append(
                (1, 1, 1, 1024, 1024, 128, is_causal, 0.2, True, None, None, False)
            )

    def set_more_shapes(self):
        return None


def flash_attention_forward_input_fn(config, dtype, device):
    (
        batch,
        num_head,
        num_head_k,
        q_seq_len,
        kv_seq_len,
        head_size,
        is_causal,
        dropout_p,
        return_debug_mask,
        window_size_left,
        window_size_right,
        use_alibi,
    ) = config

    q = torch.empty(
        (batch, q_seq_len, num_head, head_size), device=device, dtype=dtype
    ).uniform_(-0.05, 0.05)
    k = torch.empty(
        (batch, kv_seq_len, num_head_k, head_size), device=device, dtype=dtype
    ).uniform_(-0.05, 0.05)
    v = torch.empty(
        (batch, kv_seq_len, num_head_k, head_size), device=device, dtype=dtype
    ).uniform_(-0.05, 0.05)
    scale = float(1.0 / math.sqrt(head_size))

    extra_kwargs = {}
    if window_size_left is not None or window_size_right is not None:
        extra_kwargs.update(
            {
                "window_size_left": window_size_left,
                "window_size_right": window_size_right,
            }
        )
    if use_alibi:
        extra_kwargs["alibi_slopes"] = (
            torch.ones(batch, num_head, device=device, dtype=torch.float32) * 0.3
        )

    yield q, k, v, scale, is_causal, dropout_p, return_debug_mask, extra_kwargs


@pytest.mark.skipif(SkipVersion("torch", "<2.4"), reason="Low Pytorch Version.")
@pytest.mark.skipif(flag_gems.device !="cpu" and not torch.cuda.is_available(), reason="CUDA is not available")
@pytest.mark.flash_attention_forward
def test_flash_attention_forward():
    bench = FlashAttentionForwardBenchmark(
        op_name="flash_attention_forward",
        input_fn=flash_attention_forward_input_fn,
        torch_op=(gems_flash_attention_forward if flag_gems.device != "cpu" else gems_flash_attention_forward),
        dtypes=[torch.float16, torch.bfloat16],
    )
    bench.set_gems(gems_flash_attention_forward)
    bench.run()