# Copyright 2024 the HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""PyTorch Mllama model."""
import math
from typing import (List, Optional, Tuple)

import torch
import torch.nn.functional as F
import torch.utils.checkpoint
from torch import nn

from vllm import _mlu_ops as mlu_ops
from vllm._mlu_ops import flash_attention
from vllm.attention import AttentionMetadata
from vllm.attention.backends.mlu_attn import MLUFlashAttentionMetadata
from vllm.logger import init_logger
from vllm.attention.ops.paged_attn import PagedAttention
from vllm.model_executor.models.llama import LlamaDecoderLayer
from vllm.model_executor.models.mllama import (MllamaCrossAttentionDecoderLayer,
                                               MllamaTextCrossAttention,
                                               MllamaTextModel,
                                               MllamaVisionSdpaAttention)
from vllm_mlu._mlu_utils import USE_PAGED, BlockSizeInfo
from vllm_mlu.mlu_hijack_utils import MluHijackObject

logger = init_logger(__name__)


def vllm__model_executor__models__mllama__MllamaTextModel__forward(
    self,
    input_ids: torch.LongTensor,
    positions: Optional[torch.LongTensor],
    cross_attention_states: Optional[torch.LongTensor],
    cross_attention_mask: Optional[torch.LongTensor],
    kv_range_for_decode: Optional[List[Tuple[int, int]]],
    full_text_row_masked_out_mask: Optional[Tuple[torch.Tensor,
                                                  torch.Tensor]],
    kv_caches: List[torch.Tensor],
    attn_metadata: AttentionMetadata,
    skip_cross_attention: bool,
) -> torch.Tensor:
    inputs_embeds = self.embed_tokens(input_ids)
    hidden_states = inputs_embeds

    for idx, decoder_layer in enumerate(self.layers):
        if isinstance(decoder_layer, MllamaCrossAttentionDecoderLayer):
            if not skip_cross_attention:
                hidden_states = decoder_layer(
                    hidden_states=hidden_states,
                    cross_attention_states=cross_attention_states,
                    cross_attention_mask=cross_attention_mask,
                    kv_range_for_decode=kv_range_for_decode,
                    full_text_row_masked_out_mask=
                    full_text_row_masked_out_mask,
                    kv_cache=kv_caches[idx],
                    attn_metadata=attn_metadata,
                )
        elif isinstance(decoder_layer, LlamaDecoderLayer):
            '''
            =============================
            Modify by vllm_mlu
            =============================
            @brief: fuse residual into decoder layer.
            '''
            hidden_states = decoder_layer(
                positions=positions,
                hidden_states=hidden_states,
                kv_cache=kv_caches[idx],
                attn_metadata=attn_metadata,
            )
            '''
            ==================
            End of MLU Hijack
            ==================
            '''
        else:
            raise ValueError(
                f"Unknown decoder layer type {type(decoder_layer)}")
    hidden_states = self.norm(hidden_states)
    return hidden_states


def vllm__model_executor__models__mllama__MllamaVisionSdpaAttention__forward(
    self,
    hidden_state: torch.Tensor,
    attention_mask: Optional[torch.Tensor] = None,
) -> torch.Tensor:
    qkv, _ = self.qkv_proj(hidden_state)
    q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
    q = q.view(q.shape[0], q.shape[1], self.num_local_heads,
               self.head_dim)
    k = k.view(k.shape[0], k.shape[1], self.num_local_heads,
               self.head_dim)
    v = v.view(v.shape[0], v.shape[1], self.num_local_heads,
               self.head_dim)
    batch, seq_len_q, q_head_num, head_size = q.shape
    seq_len_k = k.shape[1]
    softmax_scale = head_size ** -0.5
    attention_mask = attention_mask.repeat(1, q_head_num, 1, 1)

    '''
    =============================
    Modify by vllm_mlu
    =============================
    @brief: replace SDPA with flash attn.
    '''
    attn_output = flash_attention(q, k, v,
                                  None, # out
                                  None, # cu_seq_lens_q
                                  None, # cu_seq_lens_kv
                                  None, # alibi_slop
                                  attention_mask, # attn_bias
                                  seq_len_q, # max_seq_len_q
                                  seq_len_k, # max_seq_len_kv
                                  softmax_scale, # softmax_scale
                                  False, # is_casual
                                  )
    '''
    ==================
    End of MLU Hijack
    ==================
    '''

    attn_output = attn_output.reshape(attn_output.shape[0],
                                      attn_output.shape[1], -1).contiguous()
    output, _ = self.o_proj(attn_output)
    return output

def vllm__model_executor__models__mllama__MllamaTextCrossAttention___attention_with_mask(
    self,
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    kv_cache: torch.Tensor,
    attention_mask: torch.Tensor,
    kv_range_for_decode: List[Tuple[int, int]],
    attn_metadata: AttentionMetadata,
) -> torch.Tensor:
    # Skip writing kv-cache for the initial profiling run.
    if len(kv_cache[0].shape) > 1:
        if isinstance(attn_metadata, MLUFlashAttentionMetadata):
            cached_k = torch.cat([k[s:e] for s, e in kv_range_for_decode])
            cached_v = torch.cat([v[s:e] for s, e in kv_range_for_decode])
            if USE_PAGED:
                mlu_ops.reshape_paged_cache(cached_k,
                                            cached_v,
                                            kv_cache[0][0],
                                            kv_cache[0][1],
                                            attn_metadata.cross_slot_mapping,
                        )
            else:
                cross_slot_mapping_flat = attn_metadata.cross_slot_mapping.flatten()
                seq_start_loc = attn_metadata._cached_prefill_metadata.encoder_seq_start_loc
                batch_ids = cross_slot_mapping_flat[seq_start_loc[:-1]] // BlockSizeInfo.BLOCK_SIZE
                max_context_len = attn_metadata._cached_prefill_metadata.max_encoder_seq_len
                mlu_ops.reshape_linear_cache(cached_k,
                                             cached_v,
                                             kv_cache[0][0],
                                             kv_cache[0][1],
                                             seq_start_loc, # context_lengths
                                             max_context_len,
                                             True, # packed
                                             None, # context_seq_offset
                                             batch_ids, # cache_bs_id
                                             None, # cache_seqlen_offset
                                             )
        else:
            raise ValueError(
                f"Unsupported AttentionMetadata {type(attn_metadata)} "
                f"class found. Expected the AttentionMetadata to "
                f"be either FlashAttentionMetadata.")

    '''
    =============================
    Modify by vllm_mlu
    =============================
    @brief: replace SDPA with flash attn.
    '''
    # We have to call torch.sdpa for prefill when using a
    # custom cross-attention mask. Because the mask is not a
    # standard causal mask, neither a block diagonal mask which
    # can be optimized by xformers.BlockDiagonalMask.
    # The mask is specially calculated for supporting multi
    # images and interleaved images.
    seq_len_q, q_head_num, head_size = q.shape
    softmax_scale = head_size ** -0.5
    cu_seq_lens_q = attn_metadata.seq_start_loc
    cu_seq_lens_kv = attn_metadata.encoder_seq_start_loc

    max_seq_len_q = attn_metadata.max_prefill_seq_len
    max_seq_len_kv = attn_metadata.max_encoder_seq_len
    attn_output = flash_attention(q, k, v,
                                  None, # out
                                  cu_seq_lens_q,
                                  cu_seq_lens_kv,
                                  None, # alibi_slope
                                  None, # attn_bias
                                  max_seq_len_q,
                                  max_seq_len_kv,
                                  softmax_scale,
                                  False, # is_causal
                                  )
    '''
    ==================
    End of MLU Hijack
    ==================
    '''
    output = attn_output.reshape(seq_len_q, self.num_local_heads * self.head_dim)
    return output

MluHijackObject.apply_hijack(MllamaTextCrossAttention,
                             MllamaTextCrossAttention._attention_with_mask,
                             vllm__model_executor__models__mllama__MllamaTextCrossAttention___attention_with_mask)
MluHijackObject.apply_hijack(MllamaTextModel,
                             MllamaTextModel.forward,
                             vllm__model_executor__models__mllama__MllamaTextModel__forward)
MluHijackObject.apply_hijack(MllamaVisionSdpaAttention,
                             MllamaVisionSdpaAttention.forward,
                             vllm__model_executor__models__mllama__MllamaVisionSdpaAttention__forward)
