from typing import Optional
import torch
import torch.nn.functional as F
from einops import rearrange, repeat
from vllm.logger import init_logger
from vllm.model_executor.layers.activation import QuickGELU
from vllm.model_executor.models.qwen2_vl import (
    Qwen2VisionMLP, Qwen2VisionTransformer, Qwen2VisionAttention, Qwen2VLForConditionalGeneration)
from vllm.model_executor.layers.quantization import QuantizationConfig
from vllm.config import VllmConfig
from vllm.attention.selector import _Backend
from vllm_mlu.mlu_hijack_utils import MluHijackObject

logger = init_logger(__name__)

def vllm__module_executor__models__qwen2_vl__Qwen2VisionTransformer__forward(
    self,
    x: torch.Tensor,
    grid_thw: torch.Tensor
):
    # patchify
    x = x.to(device=self.device, dtype=self.dtype)
    x = self.patch_embed(x)

    # compute position embedding
    rotary_pos_emb = self.rot_pos_emb(grid_thw)

    '''
    =============================
    Modify by vllm_mlu
    =============================
    @brief: pack q & k to fit tmo.apply_rotary
    '''
    # compute cos sin for apply_rope
    cos = rotary_pos_emb.cos()
    sin = rotary_pos_emb.sin()
    cos = repeat(cos, "... d -> ... (2 d)")
    sin = repeat(sin, "... d -> ... (2 d)")
    rotary_pos_emb.cos = cos
    rotary_pos_emb.sin = sin
    '''
    ==================
    End of MLU Hijack
    ==================
    '''

    # compute cu_seqlens
    cu_seqlens = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2],
                                         grid_thw[:, 0]).cumsum(
                                             dim=0, dtype=torch.int32)
    cu_seqlens = F.pad(cu_seqlens, (1, 0), "constant", 0)

    # transformers
    x = x.unsqueeze(1)
    for blk in self.blocks:
        x = blk(x, cu_seqlens=cu_seqlens, rotary_pos_emb=rotary_pos_emb)

    # adapter
    x = self.merger(x)
    return x

def vllm__module_executor__models__qwen2_vl__Qwen2VisionAttention__forward(
    self,
    x: torch.Tensor,
    cu_seqlens: torch.Tensor,
    rotary_pos_emb: torch.Tensor
):
    # [s, b, c] --> [s, b, head * 3 * head_dim]
    x, _ = self.qkv(x)

    # [s, b, head * 3 * head_dim] --> [s, b, head, 3 * head_dim]
    new_x_shape = x.size()[:-1] + (
        self.num_attention_heads_per_partition,
        3 * self.hidden_size_per_attention_head,
    )
    x = x.view(*new_x_shape)

    '''
    =============================
    Modify by vllm_mlu
    =============================
    @brief: apply mlu_ops.apply_rotary
    '''
    # [s, b, head, 3 * head_dim] --> 3 [b, s, head, head_dim]
    batch_size = x.shape[1]
    x = rearrange(x, "s b ... -> b s ...")
    head_dim = x.shape[-1] // 3
    q, k, v = x.split([head_dim] * 3, dim=-1)

    if rotary_pos_emb is not None:
        sin = rotary_pos_emb.sin
        cos = rotary_pos_emb.cos
        from vllm import _mlu_ops as mlu_ops
        q = q.float()
        q = mlu_ops.rotary_embedding(
            q, sin, cos, None, None, False, False, False, q.shape[1]
        )
        k = k.float()
        k = mlu_ops.rotary_embedding(
            k, sin, cos, None, None, False, False, False, k.shape[1]
        )
        q = q.type_as(v)
        k = k.type_as(v)
    '''
    ==================
    End of MLU Hijack
    ==================
    '''

    if self.attn_backend == _Backend.FLASH_ATTN:
        # from vllm_flash_attn.flash_attn_interface import (
        #   flash_attn_varlen_func)
        from flash_attn import flash_attn_varlen_func

        q, k, v = (rearrange(x, "b s ... -> (b s) ...") for x in [q, k, v])

        max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
        output = flash_attn_varlen_func(q,
                                        k,
                                        v,
                                        cu_seqlens_q=cu_seqlens,
                                        cu_seqlens_k=cu_seqlens,
                                        max_seqlen_q=max_seqlen,
                                        max_seqlen_k=max_seqlen,
                                        dropout_p=0,
                                        causal=False)

        context_layer = rearrange(output,
                                  "(b s) ... -> b s ...",
                                  b=batch_size)
    elif self.attn_backend == _Backend.TORCH_SDPA:
        seq_length = q.size(1)
        q, k, v = (rearrange(x, "b s h d -> b h s d") for x in [q, k, v])
        attention_mask = torch.zeros([1, seq_length, seq_length],
                                     device=q.device,
                                     dtype=torch.bool)
        for i in range(1, len(cu_seqlens)):
            attention_mask[..., cu_seqlens[i - 1]:cu_seqlens[i],
                           cu_seqlens[i - 1]:cu_seqlens[i]] = True
        output = F.scaled_dot_product_attention(q,
                                                k,
                                                v,
                                                attention_mask,
                                                dropout_p=0.0)
        context_layer = rearrange(output, "b h s d -> b s h d ")
    elif self.attn_backend == _Backend.XFORMERS:
        from xformers import ops as xops
        from xformers.ops.fmha.attn_bias import BlockDiagonalMask

        seqlens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()
        attn_bias = BlockDiagonalMask.from_seqlens(q_seqlen=seqlens,
                                                   kv_seqlen=None)

        context_layer = xops.memory_efficient_attention_forward(
            q, k, v, attn_bias=attn_bias, p=0, scale=None)
    elif self.attn_backend == _Backend.MLU_FLASH_ATTN:
        '''
        =============================
        Modify by vllm_mlu
        =============================
        @brief: apply mlu_ops.flash_attention
        '''
        q, k, v = (rearrange(x, "b s ... -> (b s) ...") for x in [q, k, v])
        max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
        output = mlu_ops.flash_attention(q,
                                         k,
                                         v,
                                         out=None,
                                         cu_seq_lens_q=cu_seqlens,
                                         cu_seq_lens_kv=cu_seqlens,
                                         max_seq_len_q=max_seqlen,
                                         max_seq_len_kv=max_seqlen,
                                         alibi_slope=None,
                                         attn_bias=None,
                                         softmax_scale=head_dim ** -0.5,
                                         is_causal=False)
        context_layer = rearrange(output,
                                  "(b s) ... -> b s ...",
                                  b=batch_size)
        '''
        ==================
        End of MLU Hijack
        ==================
        '''
    context_layer = rearrange(context_layer,
                              "b s h d -> s b (h d)").contiguous()

    output, _ = self.proj(context_layer)
    return output

vllm__module_executor__models__qwen2_vl__Qwen2VisionAttention__init_org = Qwen2VisionAttention.__init__

def vllm__module_executor__models__qwen2_vl__Qwen2VisionAttention____init__(
    self,
    embed_dim: Optional[int] = None,
    num_heads: Optional[int] = None,
    projection_size: Optional[int] = None,
    quant_config: Optional[QuantizationConfig] = None,
    prefix: str = "",
) -> None:
    vllm__module_executor__models__qwen2_vl__Qwen2VisionAttention__init_org(
            self, embed_dim, num_heads, projection_size, quant_config, prefix)
    '''
    =============================
    Modify by vllm_mlu
    =============================
    @brief: use mlu_ops.flash_atten for better performance
    '''
    self.attn_backend = _Backend.MLU_FLASH_ATTN
    '''
    ==================
    End of MLU Hijack
    ==================
    '''

def vllm__module_executor__models__qwen2_vl___maybe_ignore_quant_config(
    self,
    quant_config: QuantizationConfig
):
    '''
    =============================
    Modify by vllm_mlu
    =============================
    @brief: quantization for vit not yet supported
    '''
    if quant_config is not None:
        logger.warning("Quantization for VisionTransformer not yet supported.")
        return None
    '''
    ==================
    End of MLU Hijack
    ==================
    '''

def vllm_module_executor__models__qwen2_vl__Qwen2VisionMLP__forward(
    self,
    x: torch.Tensor
) -> torch.Tensor:
    x_parallel, _ = self.fc1(x)
    '''
    =============================
    Modify by vllm_mlu
    =============================
    @brief: better acc than mlu_ops.active for half precision
    '''
    if x_parallel.dtype == torch.half and isinstance(self.act, QuickGELU):
        x_parallel = self.act.forward_native(x_parallel)
    else:
        x_parallel = self.act(x_parallel)
    '''
    ==================
    End of MLU Hijack
    ==================
    '''
    x, _ = self.fc2(x_parallel)
    return x

MluHijackObject.apply_hijack(Qwen2VisionTransformer,
                             Qwen2VisionTransformer.forward,
                             vllm__module_executor__models__qwen2_vl__Qwen2VisionTransformer__forward)
MluHijackObject.apply_hijack(Qwen2VisionAttention,
                             Qwen2VisionAttention.forward,
                             vllm__module_executor__models__qwen2_vl__Qwen2VisionAttention__forward)
MluHijackObject.apply_hijack(Qwen2VisionAttention,
                             Qwen2VisionAttention.__init__,
                             vllm__module_executor__models__qwen2_vl__Qwen2VisionAttention____init__)
MluHijackObject.apply_hijack(Qwen2VLForConditionalGeneration,
                             Qwen2VLForConditionalGeneration._maybe_ignore_quant_config,
                             vllm__module_executor__models__qwen2_vl___maybe_ignore_quant_config)
MluHijackObject.apply_hijack(Qwen2VisionMLP,
                             Qwen2VisionMLP.forward,
                             vllm_module_executor__models__qwen2_vl__Qwen2VisionMLP__forward)
