import torch
from typing import List, Optional, Union
from transformers import PretrainedConfig
from vllm.attention import Attention, AttentionMetadata
from vllm.config import CacheConfig
from vllm.distributed import (get_tensor_model_parallel_rank,
                              get_tensor_model_parallel_world_size)
from vllm.sequence import IntermediateTensors
from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
from vllm.model_executor.layers.layernorm import RMSNorm
from vllm.model_executor.layers.linear import (QKVParallelLinear,
                                               RowParallelLinear)
from vllm.model_executor.layers.rotary_embedding import get_rope
from vllm.model_executor.models.baichuan import (
    _get_alibi_slopes, BaiChuanAttention,
    BaiChuanDecoderLayer, BaiChuanModel)

from vllm_mlu.model_executor.models.layer_utils import (
    decoder_layer_forward_base, decoder_model_forward_base_pp,
    is_per_tensor_smoothquant, is_per_token_smoothquant,
    quant_fusion_with_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

logger = init_logger(__name__)


def vllm__module_executor__models__baichuan__BaiChuanAttention__init__(
    self,
    hidden_size: int,
    num_heads: int,
    position_embedding: str,
    rope_theta: float = 10000,
    max_position_embeddings: int = 8192,
    cache_config: Optional[CacheConfig] = None,
    quant_config: Optional[QuantizationConfig] = None,
):
    super(BaiChuanAttention, self).__init__()
    self.hidden_size = hidden_size
    tensor_model_parallel_world_size = get_tensor_model_parallel_world_size()
    self.total_num_heads = num_heads
    assert self.total_num_heads % tensor_model_parallel_world_size == 0
    self.num_heads = (self.total_num_heads //
                        tensor_model_parallel_world_size)
    self.head_dim = hidden_size // self.total_num_heads
    self.postion_embedding = position_embedding
    self.rope_theta = rope_theta
    self.max_position_embeddings = max_position_embeddings

    # pylint: disable=invalid-name
    self.W_pack = QKVParallelLinear(
        hidden_size,
        self.head_dim,
        self.total_num_heads,
        self.total_num_heads,
        bias=False,
        quant_config=quant_config,
    )
    self.o_proj = RowParallelLinear(
        self.total_num_heads * self.head_dim,
        hidden_size,
        bias=False,
        quant_config=quant_config,
    )
    # Create the alibi slopes and slice them.
    if self.postion_embedding == "ALIBI":
        tp_rank = get_tensor_model_parallel_rank()
        head_start = tp_rank * self.num_heads
        head_end = (tp_rank + 1) * self.num_heads
        alibi_slopes = _get_alibi_slopes(self.total_num_heads)
        alibi_slopes = alibi_slopes[head_start:head_end].tolist()

        scaling = self.head_dim**-0.5
        '''
        =============================
        Modify by vllm_mlu
        =============================
        @brief: add cache_config to support kv8
        '''
        self.attn = Attention(self.num_heads,
                                self.head_dim,
                                scaling,
                                alibi_slopes=alibi_slopes,
                                cache_config=cache_config,
                                quant_config=quant_config)
        '''
        ==================
        End of MLU Hijack
        ==================
        '''
    else:
        self.rotary_emb = get_rope(
            self.head_dim,
            rotary_dim=self.head_dim,
            max_position=self.max_position_embeddings,
            base=self.rope_theta,
        )
        self.scaling = self.head_dim**-0.5
        self.attn = Attention(self.num_heads,
                                self.head_dim,
                                self.scaling,
                                cache_config=cache_config,
                                quant_config=quant_config)


def vllm__module_executor__models__baichuan__BaiChuanAttention__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.W_pack(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.num_heads * self.head_dim * 2, self.num_heads * self.head_dim], dim=-1)
    if self.postion_embedding != "ALIBI":
        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.o_proj(attn_output, residual)
    '''
    ==================
    End of MLU Hijack
    ==================
    '''
    return output


def vllm__module_executor__models__baichuan__BaiChuanDecoderLayer__init__(
    self,
    config: PretrainedConfig,
    position_embedding: str,
    cache_config: Optional[CacheConfig] = None,
    quant_config: Optional[QuantizationConfig] = None
):
    super(BaiChuanDecoderLayer, self).__init__()
    self.hidden_size = config.hidden_size
    rope_theta = getattr(config, "rope_theta", 10000)
    max_position_embeddings = getattr(config, "max_position_embeddings",
                                        8192)
    self.self_attn = BaiChuanAttention(
        hidden_size=self.hidden_size,
        num_heads=config.num_attention_heads,
        position_embedding=position_embedding,
        rope_theta=rope_theta,
        max_position_embeddings=max_position_embeddings,
        cache_config=cache_config,
        quant_config=quant_config,
    )
    '''
    =============================
    Modify by vllm_mlu
    =============================
    @brief: use FeedForward instead of MLP
    '''
    self.mlp = FeedForward(hidden_size=config.hidden_size,
                            intermediate_size=config.intermediate_size,
                            hidden_act='silu',
                            up_proj_name='gate_up_proj',
                            is_gated=True,
                            down_proj_name='down_proj',
                            bias=False,
                            quant_config=quant_config)
    '''
    ==================
    End of MLU Hijack
    ==================
    '''
    self.input_layernorm = RMSNorm(config.hidden_size,
                                    eps=config.rms_norm_eps)
    self.post_attention_layernorm = RMSNorm(config.hidden_size,
                                            eps=config.rms_norm_eps)
    '''
    =============================
    Modify by vllm_mlu
    =============================
    @brief: prepare to perf per-tensor sq cases if suitable
    '''
    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.self_attn.W_pack.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__baichuan__BaiChuanDecoderLayer__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.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:
            if self.is_per_token_sq_perf_cases:
                attn_quant_scale = self.self_attn.W_pack.smooth
                mlp_quant_scale = self.mlp.gate_up_proj.smooth
            else:
                attn_quant_scale = self.self_attn.W_pack.scale_to_int
                mlp_quant_scale = self.mlp.gate_up_proj.scale_to_int

        if self.quant_fusion_attn_layernorm is None:
            self.quant_fusion_attn_layernorm = quant_fusion_with_rmsnorm(
                self.input_layernorm, attn_quant_scale,
                dynamic_quant=self.is_per_token_sq_perf_cases)
            self.quant_fusion_mlp_layernorm = quant_fusion_with_rmsnorm(
                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=positions,
                                      hidden_states=hidden_states,
                                      kv_cache=kv_cache,
                                      attn_metadata=attn_metadata,
                                      input_layernorm=attn_layernorm,
                                      self_attn=self.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__baichuan__BaiChuanModel__forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
        intermediate_tensors: Optional[IntermediateTensors],
    ) -> Union[torch.Tensor, IntermediateTensors]:
    '''
    =============================
    Modify by vllm_mlu
    =============================
    '''
    return decoder_model_forward_base_pp(input_ids=input_ids,
                                         positions=positions,
                                         kv_caches=kv_caches,
                                         attn_metadata=attn_metadata,
                                         intermediate_tensors=intermediate_tensors,
                                         layers=self.layers,
                                         start_layer=self.start_layer,
                                         end_layer=self.end_layer,
                                         get_input_embeddings=self.embed_tokens,
                                         norm=self.norm)
    '''
    ==================
    End of MLU Hijack
    ==================
    '''


MluHijackObject.apply_hijack(BaiChuanAttention,
                             BaiChuanAttention.__init__,
                             vllm__module_executor__models__baichuan__BaiChuanAttention__init__)
MluHijackObject.apply_hijack(BaiChuanAttention,
                             BaiChuanAttention.forward,
                             vllm__module_executor__models__baichuan__BaiChuanAttention__forward)
MluHijackObject.apply_hijack(BaiChuanDecoderLayer,
                             BaiChuanDecoderLayer.__init__,
                             vllm__module_executor__models__baichuan__BaiChuanDecoderLayer__init__)
MluHijackObject.apply_hijack(BaiChuanDecoderLayer,
                             BaiChuanDecoderLayer.forward,
                             vllm__module_executor__models__baichuan__BaiChuanDecoderLayer__forward)
MluHijackObject.apply_hijack(BaiChuanModel,
                             BaiChuanModel.forward,
                             vllm__module_executor__models__baichuan__BaiChuanModel__forward)
