import torch

from torch.nn import LayerNorm
from typing import Optional
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.chatglm import GLMAttention, GLMBlock
from vllm_mlu.model_executor.models.layer_utils import (
    decoder_layer_forward_base,
    is_per_tensor_smoothquant,
    is_per_token_smoothquant,
    quant_fusion_with_layernorm,
    quant_fusion_with_rmsnorm
)

logger = init_logger(__name__)


def vllm__module_executor__models__chatglm__GLMAttention__forward(
    self,
    hidden_states: torch.Tensor,
    position_ids: 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.query_key_value(hidden_states, smooth_quant_scale)
    q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
    '''
    =============================
    Modify by vllm_mlu
    =============================
    @brief: pack q & k to fit tmo.apply_rotary
    '''
    qk, _ = qkv.split([self.q_size + self.kv_size, self.kv_size], dim=-1)
    self.rotary_emb(position_ids, qk.view(-1, self.num_heads + self.num_kv_heads, self.head_dim))
    '''
    ==================
    End of MLU Hijack
    ==================
    '''
    context_layer = self.attn(
        q,
        k,
        v,
        kv_cache,
        attn_metadata,
    )
    '''
    =============================
    Modify by vllm_mlu
    =============================
    '''
    attn_output, _ = self.dense(context_layer, residual)
    '''
    ==================
    End of MLU Hijack
    ==================
    '''
    return attn_output


def vllm__module_executor__models__chatglm__GLMBlock__init__(
    self,
    config,
    cache_config: Optional[CacheConfig] = None,
    quant_config: Optional[QuantizationConfig] = None,
):
    super(GLMBlock, self).__init__()
    self.apply_residual_connection_post_layernorm = (
        config.apply_residual_connection_post_layernorm)

    self.fp32_residual_connection = config.fp32_residual_connection

    layer_norm_func = RMSNorm if config.rmsnorm else LayerNorm
    # Layernorm on the input data.
    self.input_layernorm = layer_norm_func(config.hidden_size,
                                            eps=config.layernorm_epsilon)

    # Self attention.
    self.self_attention = GLMAttention(config, cache_config, quant_config)
    self.hidden_dropout = config.hidden_dropout

    # Layernorm on the attention output
    self.post_attention_layernorm = layer_norm_func(
        config.hidden_size, eps=config.layernorm_epsilon)

    '''
    =============================
    Modify by vllm_mlu
    =============================
    @brief: 1) use FeedForward instead of MLP
            2) prepare to perf per-tensor sq cases if suitable
    '''
    # MLP
    self.mlp = FeedForward(
        hidden_size=config.hidden_size,
        intermediate_size=config.ffn_hidden_size,
        hidden_act='silu',
        up_proj_name='dense_h_to_4h',
        is_gated=True,
        down_proj_name='dense_4h_to_h',
        bias=config.add_bias_linear,
        quant_config=quant_config
    )

    self.is_per_tesnor_sq_perf_cases = (is_per_tensor_smoothquant(quant_config) and
        not self.apply_residual_connection_post_layernorm)
    self.is_per_token_sq_perf_cases = (is_per_token_smoothquant(quant_config) and
        not self.apply_residual_connection_post_layernorm)
    if self.is_per_tesnor_sq_perf_cases or self.is_per_token_sq_perf_cases:
        self.self_attention.query_key_value.quant_method.skip_quant_input = True
        self.mlp.dense_h_to_4h.quant_method.skip_quant_input = True
        self.use_rmsnorm = config.rmsnorm
        self.quant_fusion_attn_layernorm = None
        self.quant_fusion_mlp_layernorm = None
    '''
    ==================
    End of MLU Hijack
    ==================
    '''


def vllm__module_executor__models__chatglm__GLMBlock__forward(
        self,
        position_ids: 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.input_layernorm
    mlp_layernorm = self.post_attention_layernorm
    if self.is_per_tesnor_sq_perf_cases or self.is_per_token_sq_perf_cases:
        if self.quant_fusion_attn_layernorm is None:
            quant_fusion_func = (quant_fusion_with_rmsnorm if
                                 self.use_rmsnorm else quant_fusion_with_layernorm)
            if self.is_per_token_sq_perf_cases:
                attn_quant_scale = self.self_attention.query_key_value.smooth
                mlp_quant_scale = self.mlp.dense_h_to_4h.smooth
            else:
                attn_quant_scale = self.self_attention.query_key_value.scale_to_int
                mlp_quant_scale = self.mlp.dense_h_to_4h.scale_to_int

            self.quant_fusion_attn_layernorm = quant_fusion_func(
                self.input_layernorm, attn_quant_scale,
                dynamic_quant=self.is_per_token_sq_perf_cases)
            self.quant_fusion_mlp_layernorm = quant_fusion_func(
                self.post_attention_layernorm, mlp_quant_scale,
                dynamic_quant=self.is_per_token_sq_perf_cases)
        attn_layernorm = self.quant_fusion_attn_layernorm
        mlp_layernorm = self.quant_fusion_mlp_layernorm

    return decoder_layer_forward_base(positions=position_ids,
                                      hidden_states=hidden_states,
                                      kv_cache=kv_cache,
                                      attn_metadata=attn_metadata,
                                      input_layernorm=attn_layernorm,
                                      self_attn=self.self_attention,
                                      post_layernorm=mlp_layernorm,
                                      mlp=self.mlp,
                                      apply_residual_connection_post_layernorm=self.apply_residual_connection_post_layernorm,
                                      position_name='position_ids',
                                      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
    ==================
    '''


MluHijackObject.apply_hijack(GLMAttention,
                             GLMAttention.forward,
                             vllm__module_executor__models__chatglm__GLMAttention__forward)
MluHijackObject.apply_hijack(GLMBlock,
                             GLMBlock.__init__,
                             vllm__module_executor__models__chatglm__GLMBlock__init__)
MluHijackObject.apply_hijack(GLMBlock,
                             GLMBlock.forward,
                             vllm__module_executor__models__chatglm__GLMBlock__forward)
