# SPDX-License-Identifier: Apache-2.0
# Adapt from https://github.com/KwaiVGI/VMoBA/blob/main/src/vmoba.py

import random
import time
from typing import Tuple

import torch

try:
    from flash_attn import (  # Use the new flash attention function
        flash_attn_varlen_func,
    )
    from flash_attn.flash_attn_interface import (
        _flash_attn_varlen_backward,
        _flash_attn_varlen_forward,
    )
except ImportError:

    def _unsupported(*args, **kwargs):
        raise ImportError(
            "flash-attn is not installed. Please install it, e.g., `pip install flash-attn`."
        )

    _flash_attn_varlen_forward = _unsupported
    _flash_attn_varlen_backward = _unsupported
    flash_attn_varlen_func = _unsupported

from functools import lru_cache

from einops import rearrange


@lru_cache(maxsize=16)
def calc_chunks(cu_seqlen, moba_chunk_size):
    """
    Calculate chunk boundaries.

    For vision tasks we include all chunks (even the last one which might be shorter)
    so that every chunk can be selected.
    """
    batch_sizes = cu_seqlen[1:] - cu_seqlen[:-1]
    batch_num_chunk = (batch_sizes + (moba_chunk_size - 1)) // moba_chunk_size
    cu_num_chunk = torch.ones(
        batch_num_chunk.numel() + 1,
        device=cu_seqlen.device,
        dtype=batch_num_chunk.dtype,
    )
    cu_num_chunk[1:] = batch_num_chunk.cumsum(dim=0)
    num_chunk = cu_num_chunk[-1]
    chunk_sizes = torch.full(
        (num_chunk + 1,), moba_chunk_size, dtype=torch.int32, device=cu_seqlen.device
    )
    chunk_sizes[0] = 0
    batch_last_chunk_size = batch_sizes - (batch_num_chunk - 1) * moba_chunk_size
    chunk_sizes[cu_num_chunk[1:]] = batch_last_chunk_size
    cu_chunk = chunk_sizes.cumsum(dim=-1, dtype=torch.int32)
    chunk_to_batch = torch.zeros(
        (num_chunk,), dtype=torch.int32, device=cu_seqlen.device
    )
    chunk_to_batch[cu_num_chunk[1:-1]] = 1
    chunk_to_batch = chunk_to_batch.cumsum(dim=0, dtype=torch.int32)

    # Do not filter out any chunk
    filtered_chunk_indices = torch.arange(
        num_chunk, device=cu_seqlen.device, dtype=torch.int32
    )
    num_filtered_chunk = num_chunk

    return cu_chunk, filtered_chunk_indices, num_filtered_chunk, chunk_to_batch


# --- Threshold Selection Helper Functions ---


def _select_threshold_query_head(
    gate: torch.Tensor,
    valid_gate_mask: torch.Tensor,
    gate_self_chunk_mask: torch.Tensor,
    simsum_threshold: float,
) -> torch.Tensor:
    """
    Selects chunks for each <query, head> pair based on threshold.
    Normalization and sorting happen along the chunk dimension (dim=0).
    """
    C, H, S = gate.shape
    eps = 1e-6

    # LSE‐style normalization per <head, query> (across chunks)
    gate_masked = torch.where(valid_gate_mask, gate, -torch.inf)  # Use -inf for max
    gate_min_val = torch.where(valid_gate_mask, gate, torch.inf)  # Use +inf for min

    row_min = gate_min_val.amin(dim=0)  # (H, S)
    row_max = gate_masked.amax(dim=0)  # (H, S)
    denom = row_max - row_min
    denom = torch.where(
        denom <= eps, torch.ones_like(denom), denom
    )  # avoid divide‑by‑zero

    gate_norm = (gate - row_min.unsqueeze(0)) / denom.unsqueeze(0)
    gate_norm = torch.where(valid_gate_mask, gate_norm, 0.0)  # (C, H, S)

    # 1) pull out the self‐chunk’s normalized weight for each <head,seq>
    self_norm = (gate_norm * gate_self_chunk_mask).sum(dim=0)  # (H, S)

    # 2) compute how much more normalized weight we need beyond self
    total_norm_sum = gate_norm.sum(dim=0)  # (H, S)
    remain_ratio = simsum_threshold - self_norm / (total_norm_sum + eps)  # (H, S)
    remain_ratio = torch.clamp(
        remain_ratio, min=0.0
    )  # if already ≥ thresh, no extra needed

    # 3) zero out the self‐chunk in a copy, so we only sort “others”
    others_norm = gate_norm.clone()
    others_norm[gate_self_chunk_mask] = 0.0

    # 4) sort the other chunks by descending norm, per <head,seq>
    sorted_norm, sorted_idx = torch.sort(
        others_norm, descending=True, dim=0
    )  # (C, H, S)

    # 5) cumulative‑sum the sorted norms per <head,seq>
    cumsum_others = sorted_norm.cumsum(dim=0)  # (C, H, S)

    # 6) for each <head,seq>, find the smallest k where cumsum_ratio ≥ remain_ratio
    ratio = cumsum_others / (total_norm_sum.unsqueeze(0) + eps)  # (C, H, S)
    cond = ratio >= remain_ratio.unsqueeze(0)  # (C, H, S) boolean mask
    any_cond = cond.any(dim=0)  # (H, S)
    # Find the index of the first True value along dim 0. If none, use C-1.
    cutoff = torch.where(
        any_cond,
        cond.float().argmax(dim=0),
        torch.full_like(any_cond, fill_value=C - 1),
    )  # (H, S)

    # 7) build a mask in sorted order up to that cutoff
    idx_range = torch.arange(C, device=gate.device).view(-1, 1, 1)  # (C, 1, 1)
    sorted_mask = idx_range <= cutoff.unsqueeze(0)  # (C, H, S)

    # 8) scatter it back to original chunk order
    others_mask = torch.zeros_like(gate, dtype=torch.bool)
    others_mask.scatter_(0, sorted_idx, sorted_mask)

    # 9) finally, include every self‐chunk plus all selected others
    final_gate_mask = valid_gate_mask & (others_mask | gate_self_chunk_mask)

    return final_gate_mask


def _select_threshold_block(
    gate: torch.Tensor,
    valid_gate_mask: torch.Tensor,
    gate_self_chunk_mask: torch.Tensor,
    simsum_threshold: float,
) -> torch.Tensor:
    """
    Selects <query, head> pairs for each block based on threshold.
    Normalization and sorting happen across the head and sequence dimensions (dim=1, 2).
    """
    C, H, S = gate.shape
    HS = H * S
    eps = 1e-6

    # LSE‐style normalization per block (across heads and queries)
    gate_masked = torch.where(valid_gate_mask, gate, -torch.inf)  # Use -inf for max
    gate_min_val = torch.where(valid_gate_mask, gate, torch.inf)  # Use +inf for min

    block_max = gate_masked.amax(dim=(1, 2), keepdim=True)  # (C, 1, 1)
    block_min = gate_min_val.amin(dim=(1, 2), keepdim=True)  # (C, 1, 1)
    block_denom = block_max - block_min
    block_denom = torch.where(
        block_denom <= eps, torch.ones_like(block_denom), block_denom
    )  # (C, 1, 1)

    gate_norm = (gate - block_min) / block_denom  # (C, H, S)
    gate_norm = torch.where(valid_gate_mask, gate_norm, 0.0)  # (C, H, S)

    # 1) identify normalized weights of entries that *are* self-chunks (from query perspective)
    self_norm_entries = gate_norm * gate_self_chunk_mask  # (C, H, S)
    # Sum these weights *per block*
    self_norm_sum_per_block = self_norm_entries.sum(dim=(1, 2))  # (C,)

    # 2) compute how much more normalized weight each block needs beyond its self-chunk contributions
    total_norm_sum_per_block = gate_norm.sum(dim=(1, 2))  # (C,)
    remain_ratio = simsum_threshold - self_norm_sum_per_block / (
        total_norm_sum_per_block + eps
    )  # (C,)
    remain_ratio = torch.clamp(remain_ratio, min=0.0)  # (C,)

    # 3) zero out the self‐chunk entries in a copy, so we only sort “others”
    others_norm = gate_norm.clone()
    others_norm[gate_self_chunk_mask] = 0.0  # Zero out self entries

    # 4) sort the other <head, seq> pairs by descending norm, per block
    others_flat = others_norm.contiguous().view(C, HS)  # (C, H*S)
    sorted_others_flat, sorted_indices_flat = torch.sort(
        others_flat, dim=1, descending=True
    )  # (C, H*S)

    # 5) cumulative‑sum the sorted norms per block
    cumsum_others_flat = sorted_others_flat.cumsum(dim=1)  # (C, H*S)

    # 6) for each block, find the smallest k where cumsum_ratio ≥ remain_ratio
    ratio_flat = cumsum_others_flat / (
        total_norm_sum_per_block.unsqueeze(1) + eps
    )  # (C, H*S)
    cond_flat = ratio_flat >= remain_ratio.unsqueeze(1)  # (C, H*S) boolean mask
    any_cond = cond_flat.any(dim=1)  # (C,)
    # Find the index of the first True value along dim 1. If none, use HS-1.
    cutoff_flat = torch.where(
        any_cond,
        cond_flat.float().argmax(dim=1),
        torch.full_like(any_cond, fill_value=HS - 1),
    )  # (C,)

    # 7) build a mask in sorted order up to that cutoff per block
    idx_range_flat = torch.arange(HS, device=gate.device).unsqueeze(0)  # (1, H*S)
    sorted_mask_flat = idx_range_flat <= cutoff_flat.unsqueeze(1)  # (C, H*S)

    # 8) scatter it back to original <head, seq> order per block
    others_mask_flat = torch.zeros_like(others_flat, dtype=torch.bool)  # (C, H*S)
    others_mask_flat.scatter_(1, sorted_indices_flat, sorted_mask_flat)
    others_mask = others_mask_flat.view(C, H, S)  # (C, H, S)

    # 9) finally, include every self‐chunk entry plus all selected others
    final_gate_mask = valid_gate_mask & (others_mask | gate_self_chunk_mask)

    return final_gate_mask


def _select_threshold_overall(
    gate: torch.Tensor,
    valid_gate_mask: torch.Tensor,
    gate_self_chunk_mask: torch.Tensor,
    simsum_threshold: float,
) -> torch.Tensor:
    """
    Selects <chunk, query, head> triplets globally based on threshold.
    Normalization and sorting happen across all valid entries.
    """
    C, H, S = gate.shape
    CHS = C * H * S
    eps = 1e-6

    # LSE‐style normalization globally across all valid entries
    gate_masked = torch.where(valid_gate_mask, gate, -torch.inf)  # Use -inf for max
    gate_min_val = torch.where(valid_gate_mask, gate, torch.inf)  # Use +inf for min

    overall_max = gate_masked.max()  # scalar
    overall_min = gate_min_val.min()  # scalar
    overall_denom = overall_max - overall_min
    overall_denom = torch.where(
        overall_denom <= eps,
        torch.tensor(1.0, device=gate.device, dtype=gate.dtype),
        overall_denom,
    )

    gate_norm = (gate - overall_min) / overall_denom  # (C, H, S)
    gate_norm = torch.where(valid_gate_mask, gate_norm, 0.0)  # (C, H, S)

    # 1) identify normalized weights of entries that *are* self-chunks
    self_norm_entries = gate_norm * gate_self_chunk_mask  # (C, H, S)
    # Sum these weights globally
    self_norm_sum_overall = self_norm_entries.sum()  # scalar

    # 2) compute how much more normalized weight is needed globally beyond self-chunk contributions
    total_norm_sum_overall = gate_norm.sum()  # scalar
    remain_ratio = simsum_threshold - self_norm_sum_overall / (
        total_norm_sum_overall + eps
    )  # scalar
    remain_ratio = torch.clamp(remain_ratio, min=0.0)  # scalar

    # 3) zero out the self‐chunk entries in a copy, so we only sort “others”
    others_norm = gate_norm.clone()
    others_norm[gate_self_chunk_mask] = 0.0  # Zero out self entries

    # 4) sort all other entries by descending norm, globally
    others_flat = others_norm.flatten()  # (C*H*S,)
    valid_others_mask_flat = (
        valid_gate_mask.flatten() & ~gate_self_chunk_mask.flatten()
    )  # Mask for valid, non-self entries

    # Only sort the valid 'other' entries
    valid_others_indices = torch.where(valid_others_mask_flat)[0]
    valid_others_values = others_flat[valid_others_indices]

    sorted_others_values, sort_perm = torch.sort(
        valid_others_values, descending=True
    )  # (N_valid_others,)
    sorted_original_indices = valid_others_indices[
        sort_perm
    ]  # Original indices in C*H*S space, sorted by value

    # 5) cumulative‑sum the sorted valid 'other' norms globally
    cumsum_others_values = sorted_others_values.cumsum(dim=0)  # (N_valid_others,)

    # 6) find the smallest k where cumsum_ratio ≥ remain_ratio globally
    ratio_values = cumsum_others_values / (
        total_norm_sum_overall + eps
    )  # (N_valid_others,)
    cond_values = ratio_values >= remain_ratio  # (N_valid_others,) boolean mask
    any_cond = cond_values.any()  # scalar

    # Find the index of the first True value in the *sorted* list. If none, use all valid others.
    cutoff_idx_in_sorted = torch.where(
        any_cond,
        cond_values.float().argmax(dim=0),
        torch.tensor(
            len(sorted_others_values) - 1, device=gate.device, dtype=torch.long
        ),
    )

    # 7) build a mask selecting the top-k others based on the cutoff
    # Select the original indices corresponding to the top entries in the sorted list
    selected_other_indices = sorted_original_indices[: cutoff_idx_in_sorted + 1]

    # 8) create the mask in the original flat shape
    others_mask_flat = torch.zeros_like(others_flat, dtype=torch.bool)  # (C*H*S,)
    if selected_other_indices.numel() > 0:  # Check if any 'other' indices were selected
        others_mask_flat[selected_other_indices] = True
    others_mask = others_mask_flat.view(C, H, S)  # (C, H, S)

    # 9) finally, include every self‐chunk entry plus all selected others
    final_gate_mask = valid_gate_mask & (others_mask | gate_self_chunk_mask)

    return final_gate_mask


def _select_threshold_head_global(
    gate: torch.Tensor,
    valid_gate_mask: torch.Tensor,
    gate_self_chunk_mask: torch.Tensor,
    simsum_threshold: float,
) -> torch.Tensor:
    """
    Selects <chunk, query> globally for each head based on threshold.
    """
    C, H, S = gate.shape
    eps = 1e-6

    # 1) LSE‐style normalization per head (across chunks and sequence dims)
    gate_masked = torch.where(valid_gate_mask, gate, -torch.inf)
    gate_min_val = torch.where(valid_gate_mask, gate, torch.inf)

    max_per_head = gate_masked.amax(dim=(0, 2), keepdim=True)  # (1, H, 1)
    min_per_head = gate_min_val.amin(dim=(0, 2), keepdim=True)  # (1, H, 1)
    denom = max_per_head - min_per_head
    denom = torch.where(denom <= eps, torch.ones_like(denom), denom)

    gate_norm = (gate - min_per_head) / denom
    gate_norm = torch.where(valid_gate_mask, gate_norm, 0.0)  # (C, H, S)

    # 2) sum normalized self‐chunk contributions per head
    self_norm_sum = (gate_norm * gate_self_chunk_mask).sum(dim=(0, 2))  # (H,)

    # 3) total normalized sum per head
    total_norm_sum = gate_norm.sum(dim=(0, 2))  # (H,)

    # 4) how much more normalized weight needed per head
    remain_ratio = simsum_threshold - self_norm_sum / (total_norm_sum + eps)  # (H,)
    remain_ratio = torch.clamp(remain_ratio, min=0.0)

    # 5) zero out self‐chunk entries to focus on "others"
    others_norm = gate_norm.clone()
    others_norm[gate_self_chunk_mask] = 0.0  # (C, H, S)

    # 6) flatten chunk and sequence dims, per head
    CS = C * S
    others_flat = others_norm.permute(1, 0, 2).reshape(H, CS)  # (H, C*S)
    valid_flat = (
        (valid_gate_mask & ~gate_self_chunk_mask).permute(1, 0, 2).reshape(H, CS)
    )  # (H, C*S)

    # 7) vectorized selection of “others” per head
    masked_flat = torch.where(valid_flat, others_flat, torch.zeros_like(others_flat))
    sorted_vals, sorted_idx = torch.sort(
        masked_flat, dim=1, descending=True
    )  # (H, C*S)

    cumsum_vals = sorted_vals.cumsum(dim=1)  # (H, C*S)
    ratio_vals = cumsum_vals / (total_norm_sum.unsqueeze(1) + eps)  # (H, C*S)
    cond = ratio_vals >= remain_ratio.unsqueeze(1)  # (H, C*S)

    has_cutoff = cond.any(dim=1)  # (H,)
    default = torch.full((H,), CS - 1, device=gate.device, dtype=torch.long)
    cutoff = torch.where(has_cutoff, cond.float().argmax(dim=1), default)  # (H,)

    idx_range = torch.arange(CS, device=gate.device).unsqueeze(0)  # (1, C*S)
    sorted_mask = idx_range <= cutoff.unsqueeze(1)  # (H, C*S)

    selected_flat = torch.zeros_like(valid_flat)  # (H, C*S)
    selected_flat.scatter_(1, sorted_idx, sorted_mask)  # (H, C*S)

    # 8) reshape selection mask back to (C, H, S)
    others_mask = selected_flat.reshape(H, C, S).permute(1, 0, 2)  # (C, H, S)

    # 9) include self‐chunks plus selected others, and obey valid mask
    final_gate_mask = valid_gate_mask & (gate_self_chunk_mask | others_mask)

    return final_gate_mask


class MixedAttention(torch.autograd.Function):
    @staticmethod
    def forward(
        ctx,
        q,
        k,
        v,
        self_attn_cu_seqlen,
        moba_q,
        moba_kv,
        moba_cu_seqlen_q,
        moba_cu_seqlen_kv,
        max_seqlen,
        moba_chunk_size,
        moba_q_sh_indices,
    ):
        ctx.max_seqlen = max_seqlen
        ctx.moba_chunk_size = moba_chunk_size
        ctx.softmax_scale = softmax_scale = q.shape[-1] ** (-0.5)

        # Non-causal self-attention branch
        # return out, softmax_lse, S_dmask, rng_state
        self_attn_out_sh, self_attn_lse_hs, _, _ = _flash_attn_varlen_forward(
            q=q,
            k=k,
            v=v,
            cu_seqlens_q=self_attn_cu_seqlen,
            cu_seqlens_k=self_attn_cu_seqlen,
            max_seqlen_q=max_seqlen,
            max_seqlen_k=max_seqlen,
            softmax_scale=softmax_scale,
            causal=False,
            dropout_p=0.0,
        )
        # MOBA attention branch (non-causal)
        moba_attn_out, moba_attn_lse_hs, _, _ = _flash_attn_varlen_forward(
            q=moba_q,
            k=moba_kv[:, 0],
            v=moba_kv[:, 1],
            cu_seqlens_q=moba_cu_seqlen_q,
            cu_seqlens_k=moba_cu_seqlen_kv,
            max_seqlen_q=max_seqlen,
            max_seqlen_k=moba_chunk_size,
            softmax_scale=softmax_scale,
            causal=False,
            dropout_p=0.0,
        )

        self_attn_lse_sh = self_attn_lse_hs.t().contiguous()
        moba_attn_lse = moba_attn_lse_hs.t().contiguous()

        output = torch.zeros(
            (q.shape[0], q.shape[1], q.shape[2]), device=q.device, dtype=torch.float32
        )
        output_2d = output.view(-1, q.shape[2])

        max_lse_1d = self_attn_lse_sh.view(-1)
        max_lse_1d = max_lse_1d.index_reduce(
            0, moba_q_sh_indices, moba_attn_lse.view(-1), "amax"
        )
        self_attn_lse_sh = self_attn_lse_sh - max_lse_1d.view_as(self_attn_lse_sh)
        moba_attn_lse = (
            moba_attn_lse.view(-1)
            .sub(max_lse_1d.index_select(0, moba_q_sh_indices))
            .reshape_as(moba_attn_lse)
        )

        mixed_attn_se_sh = self_attn_lse_sh.exp()
        moba_attn_se = moba_attn_lse.exp()

        mixed_attn_se_sh.view(-1).index_add_(
            0, moba_q_sh_indices, moba_attn_se.view(-1)
        )
        mixed_attn_lse_sh = mixed_attn_se_sh.log()

        # Combine self-attention output
        factor = (self_attn_lse_sh - mixed_attn_lse_sh).exp()  # [S, H]
        self_attn_out_sh = self_attn_out_sh * factor.unsqueeze(-1)
        output_2d += self_attn_out_sh.reshape_as(output_2d)

        # Combine MOBA attention output
        mixed_attn_lse = (
            mixed_attn_lse_sh.view(-1)
            .index_select(0, moba_q_sh_indices)
            .view_as(moba_attn_lse)
        )
        factor = (moba_attn_lse - mixed_attn_lse).exp()  # [S, H]
        moba_attn_out = moba_attn_out * factor.unsqueeze(-1)
        raw_attn_out = moba_attn_out.view(-1, moba_attn_out.shape[-1])
        output_2d.index_add_(0, moba_q_sh_indices, raw_attn_out)
        output = output.to(q.dtype)
        mixed_attn_lse_sh = mixed_attn_lse_sh + max_lse_1d.view_as(mixed_attn_se_sh)
        ctx.save_for_backward(
            output,
            mixed_attn_lse_sh,
            q,
            k,
            v,
            self_attn_cu_seqlen,
            moba_q,
            moba_kv,
            moba_cu_seqlen_q,
            moba_cu_seqlen_kv,
            moba_q_sh_indices,
        )

        return output

    @staticmethod
    def backward(ctx, d_output):

        max_seqlen = ctx.max_seqlen
        moba_chunk_size = ctx.moba_chunk_size
        softmax_scale = ctx.softmax_scale

        (
            output,
            mixed_attn_vlse_sh,
            q,
            k,
            v,
            self_attn_cu_seqlen,
            moba_q,
            moba_kv,
            moba_cu_seqlen_q,
            moba_cu_seqlen_kv,
            moba_q_sh_indices,
        ) = ctx.saved_tensors

        d_output = d_output.contiguous()

        dq = torch.empty_like(q)
        dk = torch.empty_like(k)
        dv = torch.empty_like(v)
        _ = _flash_attn_varlen_backward(
            dout=d_output,
            q=q,
            k=k,
            v=v,
            out=output,
            softmax_lse=mixed_attn_vlse_sh.t().contiguous(),
            dq=dq,
            dk=dk,
            dv=dv,
            cu_seqlens_q=self_attn_cu_seqlen,
            cu_seqlens_k=self_attn_cu_seqlen,
            max_seqlen_q=max_seqlen,
            max_seqlen_k=max_seqlen,
            softmax_scale=softmax_scale,
            causal=False,
            dropout_p=0.0,
            softcap=0.0,
            alibi_slopes=None,
            deterministic=True,
            window_size_left=-1,
            window_size_right=-1,
        )

        headdim = q.shape[-1]
        d_moba_output = (
            d_output.view(-1, headdim).index_select(0, moba_q_sh_indices).unsqueeze(1)
        )
        moba_output = (
            output.view(-1, headdim).index_select(0, moba_q_sh_indices).unsqueeze(1)
        )

        mixed_attn_vlse = (
            mixed_attn_vlse_sh.view(-1).index_select(0, moba_q_sh_indices).view(1, -1)
        )

        dmq = torch.empty_like(moba_q)
        dmkv = torch.empty_like(moba_kv)
        _ = _flash_attn_varlen_backward(
            dout=d_moba_output,
            q=moba_q,
            k=moba_kv[:, 0],
            v=moba_kv[:, 1],
            out=moba_output,
            softmax_lse=mixed_attn_vlse,
            dq=dmq,
            dk=dmkv[:, 0],
            dv=dmkv[:, 1],
            cu_seqlens_q=moba_cu_seqlen_q,
            cu_seqlens_k=moba_cu_seqlen_kv,
            max_seqlen_q=max_seqlen,
            max_seqlen_k=moba_chunk_size,
            softmax_scale=softmax_scale,
            causal=False,
            dropout_p=0.0,
            softcap=0.0,
            alibi_slopes=None,
            deterministic=True,
            window_size_left=-1,
            window_size_right=-1,
        )

        return dq, dk, dv, None, dmq, dmkv, None, None, None, None, None


def moba_attn_varlen(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    cu_seqlens: torch.Tensor,
    max_seqlen: int,
    moba_chunk_size: int,
    moba_topk: int,
    select_mode: str = "threshold",  # "topk" or "threshold"
    simsum_threshold: float = 0.25,
    threshold_type: str = "query_head",
) -> torch.Tensor:
    """
    Accelerated MOBA attention for vision tasks with proper LSE normalization.

    This version:
      - Splits KV into chunks.
      - For each query head, selects the top-k relevant KV chunks (including the self chunk)
        by amplifying the diagonal (self-chunk) logits.
      - Aggregates the attention outputs from the selected chunks using a log-sum-exp
        reduction so that attending to each query over the selected chunks is equivalent
        to the original algorithm.
    """
    # Stack keys and values.
    kv = torch.stack((k, v), dim=1)
    seqlen, num_head, head_dim = q.shape

    # Compute chunk boundaries.
    cu_chunk, filtered_chunk_indices, num_filtered_chunk, chunk_to_batch = calc_chunks(
        cu_seqlens, moba_chunk_size
    )

    self_attn_cu_seqlen = cu_chunk

    # Update top-k selection to include the self chunk.
    moba_topk = min(moba_topk, num_filtered_chunk)

    # --- Build filtered KV from chunks ---
    chunk_starts = cu_chunk[filtered_chunk_indices]  # [num_filtered_chunk]
    chunk_ends = cu_chunk[filtered_chunk_indices + 1]  # [num_filtered_chunk]
    chunk_lengths = chunk_ends - chunk_starts  # [num_filtered_chunk]
    max_chunk_len = int(chunk_lengths.max().item())

    range_tensor = torch.arange(
        max_chunk_len, device=kv.device, dtype=chunk_starts.dtype
    ).unsqueeze(0)
    indices = chunk_starts.unsqueeze(1) + range_tensor
    indices = torch.clamp(indices, max=kv.shape[0] - 1)
    valid_mask = range_tensor < chunk_lengths.unsqueeze(1)
    gathered = kv[indices.view(-1)].view(
        num_filtered_chunk, max_chunk_len, *kv.shape[1:]
    )
    gathered = gathered * valid_mask.unsqueeze(-1).unsqueeze(-1).unsqueeze(-1).type_as(
        gathered
    )

    # Compute key_gate_weight over valid tokens.
    key_values = gathered[
        :, :, 0
    ].float()  # [num_filtered_chunk, max_chunk_len, num_head, head_dim]
    valid_mask_exp = valid_mask.unsqueeze(-1).unsqueeze(-1)
    key_sum = (key_values * valid_mask_exp).sum(dim=1)
    divisor = valid_mask.sum(dim=1).unsqueeze(-1).unsqueeze(-1)
    key_gate_weight = key_sum / divisor  # [num_filtered_chunk, num_head, head_dim]

    # Compute gate logits between key_gate_weight and queries.
    q_float = q.float()
    # gate = torch.einsum("nhd,shd->nhs", key_gate_weight, q_float)  # [num_filtered_chunk, num_head, seqlen]
    gate = torch.bmm(
        key_gate_weight.permute(1, 0, 2), q_float.permute(1, 0, 2).transpose(1, 2)
    ).permute(1, 0, 2)

    # Amplify the diagonal (self chunk) contributions.
    gate_seq_idx = (
        torch.arange(seqlen, device=q.device, dtype=torch.int32)
        .unsqueeze(0)
        .expand(num_filtered_chunk, seqlen)
    )
    chunk_start = cu_chunk[filtered_chunk_indices]  # [num_filtered_chunk]
    chunk_end = cu_chunk[filtered_chunk_indices + 1]  # [num_filtered_chunk]
    gate_self_chunk_mask = (
        (
            (gate_seq_idx >= chunk_start.unsqueeze(1))
            & (gate_seq_idx < chunk_end.unsqueeze(1))
        )
        .unsqueeze(1)
        .expand(-1, num_head, -1)
    )
    amplification_factor = 1e9  # Example factor; adjust as needed.
    origin_gate = gate.clone()
    gate = gate.clone()
    if select_mode == "topk":
        gate[gate_self_chunk_mask] += amplification_factor

    # Exclude positions that are outside the valid batch boundaries.
    batch_starts = cu_seqlens[chunk_to_batch[filtered_chunk_indices]]
    batch_ends = cu_seqlens[chunk_to_batch[filtered_chunk_indices] + 1]
    gate_batch_start_mask = gate_seq_idx < batch_starts.unsqueeze(1)
    gate_batch_end_mask = gate_seq_idx >= batch_ends.unsqueeze(1)
    gate_inf_mask = gate_batch_start_mask | gate_batch_end_mask
    gate.masked_fill_(gate_inf_mask.unsqueeze(1), -float("inf"))

    if select_mode == "topk":
        # We amplify self‐chunk in gate already, so self entries will rank highest.
        valid_gate_mask = gate != -float("inf")
        if threshold_type == "query_head":
            # === per‐<head,seq> top-k across chunks (original behavior) ===
            # gate: (C, H, S)
            _, gate_topk_idx = torch.topk(
                gate, k=moba_topk, dim=0, largest=True, sorted=False
            )
            gate_idx_mask = torch.zeros_like(gate, dtype=torch.bool)
            gate_idx_mask.scatter_(0, gate_topk_idx, True)
            gate_mask = valid_gate_mask & gate_idx_mask
        elif threshold_type == "overall":
            # === global top-k across all (chunk, head, seq) entries ===
            C, H, S = gate.shape
            flat_gate = gate.flatten()
            flat_mask = valid_gate_mask.flatten()
            flat_gate_masked = torch.where(flat_mask, flat_gate, -float("inf"))
            # pick topk global entries
            vals, idx = torch.topk(
                flat_gate_masked, k=moba_topk * H * S, largest=True, sorted=False
            )
            others_mask_flat = torch.zeros_like(flat_mask, dtype=torch.bool)
            others_mask_flat[idx] = True
            gate_mask = (valid_gate_mask.flatten() & others_mask_flat).view(gate.shape)
        elif threshold_type == "head_global":
            # per-head top-k across all chunks and sequence positions
            C, H, S = gate.shape
            CS = C * S
            flat_gate = gate.permute(1, 0, 2).reshape(H, CS)
            flat_valid = valid_gate_mask.permute(1, 0, 2).reshape(H, CS)
            flat_gate_masked = torch.where(
                flat_valid, flat_gate, torch.full_like(flat_gate, -float("inf"))
            )
            # pick top-k indices per head
            _, topk_idx = torch.topk(
                flat_gate_masked, k=moba_topk * S, dim=1, largest=True, sorted=False
            )
            gate_idx_flat = torch.zeros_like(flat_valid, dtype=torch.bool)
            gate_idx_flat.scatter_(1, topk_idx, True)
            gate_mask = gate_idx_flat.reshape(H, C, S).permute(1, 0, 2)
        else:
            raise ValueError(
                f"Invalid threshold_type for topk: {threshold_type}. "
                "Choose 'query_head', 'block', or 'overall'."
            )
    elif select_mode == "threshold":
        # Delegate to the specific thresholding function
        valid_gate_mask = gate != -float("inf")  # (num_chunk, num_head, seqlen)
        if threshold_type == "query_head":
            gate_mask = _select_threshold_query_head(
                gate, valid_gate_mask, gate_self_chunk_mask, simsum_threshold
            )
        elif threshold_type == "block":
            gate_mask = _select_threshold_block(
                gate, valid_gate_mask, gate_self_chunk_mask, simsum_threshold
            )
        elif threshold_type == "overall":
            gate_mask = _select_threshold_overall(
                gate, valid_gate_mask, gate_self_chunk_mask, simsum_threshold
            )
        elif threshold_type == "head_global":
            gate_mask = _select_threshold_head_global(
                gate, valid_gate_mask, gate_self_chunk_mask, simsum_threshold
            )
        else:
            raise ValueError(
                f"Invalid threshold_type: {threshold_type}. Choose 'query_head', 'block', or 'overall'."
            )
    else:
        raise ValueError(
            f"Invalid select_mode: {select_mode}. Choose 'topk' or 'threshold'."
        )

    # eliminate self_chunk in MoBA branch
    gate_mask = gate_mask & ~gate_self_chunk_mask
    # if gate_mask is all false, perform flash_attn instead
    if gate_mask.sum() == 0:
        return flash_attn_varlen_func(
            q, k, v, cu_seqlens, cu_seqlens, max_seqlen, max_seqlen, causal=False
        )

    # Determine which query positions are selected.
    # nonzero_indices has shape [N, 3] where each row is [chunk_index, head_index, seq_index].
    moba_q_indices = gate_mask.reshape(gate_mask.shape[0], -1).nonzero(as_tuple=True)[
        -1
    ]  # [(h s k)]
    moba_q_sh_indices = (moba_q_indices % seqlen) * num_head + (
        moba_q_indices // seqlen
    )
    moba_q = (
        rearrange(q, "s h d -> (h s) d").index_select(0, moba_q_indices).unsqueeze(1)
    )

    # Build cumulative sequence lengths for the selected queries.
    moba_seqlen_q = gate_mask.sum(dim=-1).flatten()
    q_zero_mask = moba_seqlen_q == 0
    valid_expert_mask = ~q_zero_mask
    if q_zero_mask.sum() > 0:
        moba_seqlen_q = moba_seqlen_q[valid_expert_mask]
    moba_cu_seqlen_q = torch.cat(
        (
            torch.tensor([0], device=q.device, dtype=moba_seqlen_q.dtype),
            moba_seqlen_q.cumsum(dim=0),
        ),
        dim=0,
    ).to(torch.int32)

    # Rearrange gathered KV for the MOBA branch.
    experts_tensor = rearrange(gathered, "nc cl two h d -> (nc h) cl two d")
    valid_expert_lengths = (
        chunk_lengths.unsqueeze(1)
        .expand(num_filtered_chunk, num_head)
        .reshape(-1)
        .to(torch.int32)
    )
    if q_zero_mask.sum() > 0:
        experts_tensor = experts_tensor[valid_expert_mask]
        valid_expert_lengths = valid_expert_lengths[valid_expert_mask]

    seq_range = torch.arange(
        experts_tensor.shape[1], device=experts_tensor.device
    ).unsqueeze(0)
    mask = seq_range < valid_expert_lengths.unsqueeze(1)
    moba_kv = experts_tensor[mask]  # Shape: ((nc h cl_valid) two d)
    moba_kv = moba_kv.unsqueeze(2)  # Shape: ((nc h cl_valid) two 1 d)

    moba_cu_seqlen_kv = torch.cat(
        [
            torch.zeros(1, device=experts_tensor.device, dtype=torch.int32),
            valid_expert_lengths.cumsum(dim=0),
        ],
        dim=0,
    ).to(torch.int32)

    assert (
        moba_cu_seqlen_kv.shape == moba_cu_seqlen_q.shape
    ), f"Mismatch between moba_cu_seqlen_kv.shape and moba_cu_seqlen_q.shape: {moba_cu_seqlen_kv.shape} vs {moba_cu_seqlen_q.shape}"

    return MixedAttention.apply(
        q,
        k,
        v,
        self_attn_cu_seqlen,
        moba_q,
        moba_kv,
        moba_cu_seqlen_q,
        moba_cu_seqlen_kv,
        max_seqlen,
        moba_chunk_size,
        moba_q_sh_indices,
    )


def process_moba_input(
    x,
    patch_resolution,
    chunk_size,
):
    """
    Process inputs for the attention function.

    Args:
        x (torch.Tensor): Input tensor with shape [batch_size, num_patches, num_heads, head_dim].
        patch_resolution (tuple): Tuple containing the patch resolution (t, h, w).
        chunk_size (int): Size of the chunk. (maybe tuple or int, according to chunk type)

    Returns:
        torch.Tensor: Processed input tensor.
    """
    if isinstance(chunk_size, float) or isinstance(chunk_size, int):
        moba_chunk_size = int(chunk_size * patch_resolution[1] * patch_resolution[2])
    else:
        assert isinstance(
            chunk_size, (Tuple, list)
        ), f"chunk_size should be a tuple, list, or int, now it is: {type(chunk_size)}"
        if len(chunk_size) == 2:
            assert (
                patch_resolution[1] % chunk_size[0] == 0
                and patch_resolution[2] % chunk_size[1] == 0
            ), f"spatial patch_resolution {patch_resolution[1:]} should be divisible by 2d chunk_size {chunk_size}"
            nch, ncw = (
                patch_resolution[1] // chunk_size[0],
                patch_resolution[2] // chunk_size[1],
            )
            x = rearrange(
                x,
                "b (t nch ch ncw cw) n d -> b (nch ncw t ch cw) n d",
                t=patch_resolution[0],
                nch=nch,
                ncw=ncw,
                ch=chunk_size[0],
                cw=chunk_size[1],
            )
            moba_chunk_size = patch_resolution[0] * chunk_size[0] * chunk_size[1]
        elif len(chunk_size) == 3:
            assert (
                patch_resolution[0] % chunk_size[0] == 0
                and patch_resolution[1] % chunk_size[1] == 0
                and patch_resolution[2] % chunk_size[2] == 0
            ), f"patch_resolution {patch_resolution} should be divisible by 3d chunk_size {chunk_size}"
            nct, nch, ncw = (
                patch_resolution[0] // chunk_size[0],
                patch_resolution[1] // chunk_size[1],
                patch_resolution[2] // chunk_size[2],
            )
            x = rearrange(
                x,
                "b (nct ct nch ch ncw cw) n d -> b (nct nch ncw ct ch cw) n d",
                nct=nct,
                nch=nch,
                ncw=ncw,
                ct=chunk_size[0],
                ch=chunk_size[1],
                cw=chunk_size[2],
            )
            moba_chunk_size = chunk_size[0] * chunk_size[1] * chunk_size[2]
        else:
            raise ValueError(
                f"chunk_size should be a int, or a tuple of length 2 or 3, now it is: {len(chunk_size)}"
            )

    return x, moba_chunk_size


def process_moba_output(
    x,
    patch_resolution,
    chunk_size,
):
    if isinstance(chunk_size, float) or isinstance(chunk_size, int):
        pass
    elif len(chunk_size) == 2:
        x = rearrange(
            x,
            "b (nch ncw t ch cw) n d -> b (t nch ch ncw cw) n d",
            nch=patch_resolution[1] // chunk_size[0],
            ncw=patch_resolution[2] // chunk_size[1],
            t=patch_resolution[0],
            ch=chunk_size[0],
            cw=chunk_size[1],
        )
    elif len(chunk_size) == 3:
        x = rearrange(
            x,
            "b (nct nch ncw ct ch cw) n d -> b (nct ct nch ch ncw cw) n d",
            nct=patch_resolution[0] // chunk_size[0],
            nch=patch_resolution[1] // chunk_size[1],
            ncw=patch_resolution[2] // chunk_size[2],
            ct=chunk_size[0],
            ch=chunk_size[1],
            cw=chunk_size[2],
        )

    return x


# TEST
def generate_data(batch_size, seqlen, num_head, head_dim, dtype):
    random.seed(0)
    torch.manual_seed(0)
    torch.cuda.manual_seed(0)
    device = torch.cuda.current_device()

    q = torch.randn((batch_size, seqlen, num_head, head_dim), requires_grad=True).to(
        dtype=dtype, device="cuda"
    )
    k = torch.randn((batch_size, seqlen, num_head, head_dim), requires_grad=True).to(
        dtype=dtype, device="cuda"
    )
    v = torch.randn((batch_size, seqlen, num_head, head_dim), requires_grad=True).to(
        dtype=dtype, device="cuda"
    )
    print(f"q.shape: {q.shape}, k.shape: {k.shape}, v.shape: {v.shape}")
    cu_seqlens = torch.arange(
        0, q.shape[0] * q.shape[1] + 1, q.shape[1], dtype=torch.int32, device="cuda"
    )
    max_seqlen = q.shape[1]
    q = rearrange(q, "b s ... -> (b s) ...")
    k = rearrange(k, "b s ... -> (b s) ...")
    v = rearrange(v, "b s ... -> (b s) ...")

    return q, k, v, cu_seqlens, max_seqlen


def test_attn_varlen_moba_speed(
    batch,
    head,
    seqlen,
    head_dim,
    moba_chunk_size,
    moba_topk,
    dtype=torch.bfloat16,
    select_mode="threshold",
    simsum_threshold=0.25,
    threshold_type="query_head",
):
    """Speed test comparing flash_attn vs moba_attention"""
    # Get data
    q, k, v, cu_seqlen, max_seqlen = generate_data(batch, seqlen, head, head_dim, dtype)
    print(
        f"batch:{batch} head:{head} seqlen:{seqlen} chunk:{moba_chunk_size} topk:{moba_topk} select_mode: {select_mode} simsum_threshold:{simsum_threshold}"
    )
    vo_grad = torch.randn_like(q)

    # Warmup
    warmup_iters = 3
    perf_test_iters = 10

    # Warmup
    for _ in range(warmup_iters):
        o = flash_attn_varlen_func(
            q, k, v, cu_seqlen, cu_seqlen, max_seqlen, max_seqlen, causal=False
        )
        torch.autograd.backward(o, vo_grad)

    torch.cuda.synchronize()
    start_flash = time.perf_counter()
    for _ in range(perf_test_iters):
        o = flash_attn_varlen_func(
            q, k, v, cu_seqlen, cu_seqlen, max_seqlen, max_seqlen, causal=False
        )
        torch.autograd.backward(o, vo_grad)

    torch.cuda.synchronize()
    time_flash = (time.perf_counter() - start_flash) / perf_test_iters * 1000

    # Warmup
    for _ in range(warmup_iters):
        om = moba_attn_varlen(
            q,
            k,
            v,
            cu_seqlen,
            max_seqlen,
            moba_chunk_size=moba_chunk_size,
            moba_topk=moba_topk,
            select_mode=select_mode,
            simsum_threshold=simsum_threshold,
            threshold_type=threshold_type,
        )
        torch.autograd.backward(om, vo_grad)

    torch.cuda.synchronize()
    start_moba = time.perf_counter()
    for _ in range(perf_test_iters):
        om = moba_attn_varlen(
            q,
            k,
            v,
            cu_seqlen,
            max_seqlen,
            moba_chunk_size=moba_chunk_size,
            moba_topk=moba_topk,
            select_mode=select_mode,
            simsum_threshold=simsum_threshold,
            threshold_type=threshold_type,
        )
        torch.autograd.backward(om, vo_grad)

    torch.cuda.synchronize()
    time_moba = (time.perf_counter() - start_moba) / perf_test_iters * 1000

    print(f"Flash: {time_flash:.2f}ms, MoBA: {time_moba:.2f}ms")
    print(f"Speedup:  {time_flash / time_moba:.2f}x")


if __name__ == "__main__":
    """
    CUDA_VISIBLE_DEVICES=1 \
    python -u csrc/attn/vmoba_attn/vmoba/vmoba.py
    """
    test_attn_varlen_moba_speed(
        batch=1,
        head=12,
        seqlen=32760,
        head_dim=128,
        moba_chunk_size=32760 // 3 // 6 // 4,
        moba_topk=3,
        select_mode="threshold",
        simsum_threshold=0.3,
        threshold_type="query_head",
    )