import torch
import numpy as np

from typing import List, Optional, Union
from transformers import PretrainedConfig
from vllm.attention import AttentionMetadata
from vllm.config import CacheConfig
from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
from vllm.model_executor.layers.layernorm import RMSNorm
from vllm_mlu.model_executor.layers.feed_forward import FeedForward
from vllm_mlu.mlu_hijack_utils import MluHijackObject
from vllm.logger import init_logger

from vllm.model_executor.models.qwen import QWenAttention, QWenBlock, QWenModel, QwenImageInputs
from vllm.sequence import IntermediateTensors
from vllm.distributed import get_pp_group
from .layer_utils import decoder_layer_forward_base
from vllm_mlu.model_executor.models.layer_utils import (
    decoder_layer_forward_base, decoder_model_forward_base,
    is_per_tensor_smoothquant, is_per_token_smoothquant,
    quant_fusion_with_rmsnorm)

logger = init_logger(__name__)


def vllm__module_executor__models__qwen__QwenAttention__forward(
    self,
    positions: torch.Tensor,
    hidden_states: torch.Tensor,
    kv_cache: torch.Tensor,
    attn_metadata: AttentionMetadata,
    residual: Optional[torch.Tensor] = None,
    smooth_quant_scale: Optional[torch.Tensor] = None,
) -> torch.Tensor:
    qkv, _ = self.c_attn(hidden_states, smooth_quant_scale)
    q, k, v = qkv.chunk(chunks=3, dim=-1)
    '''
    =============================
    Modify by vllm_mlu
    =============================
    @brief: pack q & k to fit tmo.apply_rotary
    '''
    qk, _ = qkv.split([self.head_dim * self.num_heads * 2, self.head_dim * self.num_heads], dim=-1)
    self.rotary_emb(positions, qk.view(-1, self.num_heads + self.num_heads, self.head_dim))
    '''
    ==================
    End of MLU Hijack
    ==================
    '''
    attn_output = self.attn(q, k, v, kv_cache, attn_metadata)
    '''
    =============================
    Modify by vllm_mlu
    =============================
    @brief: add residual
    '''
    output, _ = self.c_proj(attn_output, residual)
    '''
    ==================
    End of MLU Hijack
    ==================
    '''
    return output


def vllm__module_executor__models__qwen__QWenBlock__init__(
    self,
    config: PretrainedConfig,
    cache_config: Optional[CacheConfig] = None,
    quant_config: Optional[QuantizationConfig] = None,
):
    super(QWenBlock, self).__init__()
    self.ln_1 = RMSNorm(config.hidden_size, eps=config.layer_norm_epsilon)

    rope_theta = getattr(config, "rope_theta", 10000)
    rope_scaling = getattr(config, "rope_scaling", None)
    self.attn = QWenAttention(config.hidden_size,
                              config.num_attention_heads,
                              config.max_position_embeddings,
                              rope_theta=rope_theta,
                              rope_scaling=rope_scaling,
                              cache_config=cache_config,
                              quant_config=quant_config)

    self.ln_2 = RMSNorm(config.hidden_size, eps=config.layer_norm_epsilon)
    '''
    =============================
    Modify by vllm_mlu
    =============================
    @brief: 1) use FeedForward instead of MLP
            2) prepare to perf per-tensor sq cases if suitable
    '''
    self.mlp = FeedForward(hidden_size=config.hidden_size,
                            intermediate_size=config.intermediate_size // 2,
                            hidden_act='silu',
                            up_proj_name='gate_up_proj',
                            is_gated=True,
                            down_proj_name='c_proj',
                            bias=False,
                            quant_config=quant_config)

    self.is_per_tesnor_sq_perf_cases = is_per_tensor_smoothquant(quant_config)
    self.is_per_token_sq_perf_cases = is_per_token_smoothquant(quant_config)
    if self.is_per_tesnor_sq_perf_cases or self.is_per_token_sq_perf_cases:
        self.attn.c_attn.quant_method.skip_quant_input = True
        self.mlp.gate_up_proj.quant_method.skip_quant_input = True
        self.quant_fusion_attn_layernorm = None
        self.quant_fusion_mlp_layernorm = None
    '''
    ==================
    End of MLU Hijack
    ==================
    '''


def vllm__module_executor__models__qwen__QWenBlock__forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
        kv_cache: torch.Tensor,
        attn_metadata: AttentionMetadata,
    ) -> torch.Tensor:
    '''
    =============================
    Modify by vllm_mlu
    =============================
    @brief: perf model by:
    1) add residual in matmul;
    2) fuse quantization in layernorm in per-tensor sq case;
    '''
    attn_layernorm = self.ln_1
    mlp_layernorm = self.ln_2
    if self.is_per_tesnor_sq_perf_cases:
        if self.quant_fusion_attn_layernorm is None:
            self.quant_fusion_attn_layernorm = quant_fusion_with_rmsnorm(
                self.ln_1, self.attn.c_attn.scale_to_int)
            self.quant_fusion_mlp_layernorm = quant_fusion_with_rmsnorm(
                self.ln_2, self.mlp.gate_up_proj.scale_to_int)
        attn_layernorm = self.quant_fusion_attn_layernorm
        mlp_layernorm = self.quant_fusion_mlp_layernorm
    elif self.is_per_token_sq_perf_cases:
        if self.quant_fusion_attn_layernorm is None:
            self.quant_fusion_attn_layernorm = quant_fusion_with_rmsnorm(
                self.ln_1, self.attn.c_attn.smooth, dynamic_quant=True)
            self.quant_fusion_mlp_layernorm = quant_fusion_with_rmsnorm(
                self.ln_2, self.mlp.gate_up_proj.smooth, dynamic_quant=True)
        attn_layernorm = self.quant_fusion_attn_layernorm
        mlp_layernorm = self.quant_fusion_mlp_layernorm

    return decoder_layer_forward_base(positions=positions,
                                      hidden_states=hidden_states,
                                      kv_cache=kv_cache,
                                      attn_metadata=attn_metadata,
                                      input_layernorm=attn_layernorm,
                                      self_attn=self.attn,
                                      post_layernorm=mlp_layernorm,
                                      mlp=self.mlp,
                                      input_norm_fuse_en=self.is_per_token_sq_perf_cases,
                                      post_norm_fuse_en=self.is_per_token_sq_perf_cases)
    '''
    ==================
    End of MLU Hijack
    ==================
    '''


def vllm__module_executor__models__qwen__QWenModel__forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
        intermediate_tensors: Optional[IntermediateTensors],
        pixel_values: Optional[QwenImageInputs],
    ) -> Union[torch.Tensor, IntermediateTensors]:
    img_pos = None
    # If pixel / visual embeddings are provided, this is a visual model
    if pixel_values is not None and self.visual is not None:
        if pixel_values["type"] != "image_embeds":
            image_embeds = self.visual(pixel_values["data"])
        else:
            image_embeds = pixel_values["data"]

        # features should be of shape (# images, 256, hidden_dim)
        img_pos = self.visual.get_image_positions(input_ids)
        if isinstance(
                img_pos,
                np.ndarray) and img_pos.shape[0] != image_embeds.shape[0]:
            raise ValueError(
                f"Number of placeholders: {img_pos.shape[0]} "
                f"does not match number of images {image_embeds.shape[0]}."
            )

    '''
    =============================
    Modify by vllm_mlu
    =============================
    @brief: remove residual
    '''
    if get_pp_group().is_first_rank:
        hidden_states = self.wte(input_ids)
        # Merge the image embeddings into the hidden states if actually have
        # visual features and the corresponding image tokens
        if img_pos is not None:
            for idx, (img_bos, img_eos) in enumerate(img_pos):
                hidden_states[img_bos + 1:img_eos] = image_embeds[idx]
    else:
        assert intermediate_tensors is not None
        hidden_states = intermediate_tensors["hidden_states"]
    for i in range(self.start_layer, self.end_layer):
        layer = self.h[i]
        hidden_states = layer(
            positions,
            hidden_states,
            kv_caches[i - self.start_layer],
            attn_metadata,
        )
    if not get_pp_group().is_last_rank:
        return IntermediateTensors({
            "hidden_states": hidden_states,
        })
    hidden_states = self.ln_f(hidden_states)
    '''
    ==================
    End of MLU Hijack
    ==================
    '''
    return hidden_states

MluHijackObject.apply_hijack(QWenAttention,
                             QWenAttention.forward,
                             vllm__module_executor__models__qwen__QwenAttention__forward)
MluHijackObject.apply_hijack(QWenBlock,
                             QWenBlock.__init__,
                             vllm__module_executor__models__qwen__QWenBlock__init__)
MluHijackObject.apply_hijack(QWenBlock,
                             QWenBlock.forward,
                             vllm__module_executor__models__qwen__QWenBlock__forward)
MluHijackObject.apply_hijack(QWenModel,
                             QWenModel.forward,
                             vllm__module_executor__models__qwen__QWenModel__forward)
