import torch
from typing import Optional, Tuple, Iterable, Union, List
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.internlm2 import InternLM2Attention, InternLMDecoderLayer, InternLM2ForCausalLM, InternLM2Model
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
from vllm.model_executor.models.utils import (is_pp_missing_parameter)
from vllm.sequence import IntermediateTensors
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)

logger = init_logger(__name__)


def vllm__module_executor__models__internlm2__InternLM2Attention__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.wqkv(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, v = qkv.split([self.q_size + self.kv_size, self.kv_size], dim=-1)
    self.rotary_emb(positions, qk.view(-1, self.num_heads + self.num_kv_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.wo(attn_output, residual)
    '''
    ==================
    End of MLU Hijack
    ==================
    '''
    return output


def vllm__module_executor__models__internlm2__InternLMDecoderLayer____init__(
    self,
    config: PretrainedConfig,
    cache_config: Optional[CacheConfig] = None,
    quant_config: Optional[QuantizationConfig] = None,
    prefix: str = "") -> None:
    super(InternLMDecoderLayer, self).__init__()
    self.hidden_size = config.hidden_size
    rope_theta = getattr(config, "rope_theta", 10000)
    rope_scaling = getattr(config, "rope_scaling", None)
    max_position_embeddings = getattr(config, "max_position_embeddings",
                                        8192)
    self.attention = InternLM2Attention(
        hidden_size=self.hidden_size,
        num_heads=config.num_attention_heads,
        num_kv_heads=config.num_key_value_heads,
        rope_theta=rope_theta,
        rope_scaling=rope_scaling,
        max_position_embeddings=max_position_embeddings,
        cache_config=cache_config,
        quant_config=quant_config,
        prefix=f"{prefix}.attention",
    )
    '''
    =============================
    Modify by vllm_mlu
    =============================
    @brief: use FeedForward instead of MLP
    '''
    self.feed_forward = FeedForward(
        hidden_size=self.hidden_size,
        intermediate_size=config.intermediate_size,
        hidden_act=config.hidden_act,
        up_proj_name='gate_up_proj',
        is_gated=True,
        down_proj_name='w2',
        bias=False,
        quant_config=quant_config,
        prefix=f'{prefix}.feed_forward'
    )
    '''
    ==================
    End of MLU Hijack
    ==================
    '''
    self.attention_norm = RMSNorm(config.hidden_size,
                                    eps=config.rms_norm_eps)
    self.ffn_norm = 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.attention.wqkv.quant_method.skip_quant_input = True
        self.feed_forward.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__internlm2__InternLM2ForCausalLM__load_weights(
                                self, weights: Iterable[Tuple[str, torch.Tensor]]):
    stacked_params_mapping = [
        # (param_name, shard_name, shard_id)
        ("gate_up_proj", "w1", 0),
        ("gate_up_proj", "w3", 1),
    ]
    params_dict = dict(self.named_parameters())
    for name, loaded_weight in weights:
        if "rotary_emb.inv_freq" in name:
            continue
        for (param_name, weight_name, shard_id) in stacked_params_mapping:
            if weight_name not in name:
                continue
            name = name.replace(weight_name, param_name)
            # Skip loading extra bias for GPTQ models.
            if name.endswith(".bias") and name not in params_dict:
                continue
            if is_pp_missing_parameter(name, self):
                continue
            param = params_dict[name]
            weight_loader = param.weight_loader
            weight_loader(param, loaded_weight, shard_id)
            break
        else:
            # Skip loading extra bias for GPTQ models.
            if name.endswith(".bias") and name not in params_dict:
                continue
            if is_pp_missing_parameter(name, self):
                continue
            param = params_dict[name]
            '''
            =============================
            Modify by vllm_mlu
            =============================
            @brief: support load quant weights and params
            '''
            if "wqkv" in name and 'smooth' not in name and 'scale_to_int' not in name:
                config = self.config
                kv_groups = (config.num_attention_heads //
                                config.num_key_value_heads)
                head_dim = config.hidden_size // config.num_attention_heads
                if 'weight' in name:
                    loaded_weight = loaded_weight.view(-1, 2 + kv_groups,
                                                        head_dim,
                                                        loaded_weight.shape[-1])
                    wq, wk, wv = torch.split(loaded_weight, [kv_groups, 1, 1],
                                                dim=1)
                    wq = wq.reshape(-1, wq.shape[-1])
                    wk = wk.reshape(-1, wk.shape[-1])
                    wv = wv.reshape(-1, wv.shape[-1])
                elif 'scale' in name:
                    loaded_weight = loaded_weight.view(-1, 2 + kv_groups, head_dim)
                    wq, wk, wv = torch.split(loaded_weight, [kv_groups, 1, 1],
                                                dim=1)
                    wq = wq.reshape(-1)
                    wk = wk.reshape(-1)
                    wv = wv.reshape(-1)
                else:
                    logger.error(f"unsupport internlm2 quant param: {name}, shape: {loaded_weight.shape}")
                weight_loader = param.weight_loader
                weight_loader(param, wq, 'q')
                weight_loader(param, wk, 'k')
                weight_loader(param, wv, 'v')
            else:
                weight_loader = getattr(param, "weight_loader",
                                        default_weight_loader)
                weight_loader(param, loaded_weight)
            '''
            ==================
            End of MLU Hijack
            ==================
            '''


def vllm__module_executor__models__internlm2__InternLMDecoderLayer__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.attention_norm
    mlp_layernorm = self.ffn_norm
    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.attention.wqkv.smooth
                mlp_quant_scale = self.feed_forward.gate_up_proj.smooth
            else:
                attn_quant_scale = self.attention.wqkv.scale_to_int
                mlp_quant_scale = self.feed_forward.gate_up_proj.scale_to_int
            self.quant_fusion_attn_layernorm = quant_fusion_with_rmsnorm(
                self.attention_norm, attn_quant_scale,
                dynamic_quant=self.is_per_token_sq_perf_cases)
            self.quant_fusion_mlp_layernorm = quant_fusion_with_rmsnorm(
                self.ffn_norm, 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, hidden_states, kv_cache, attn_metadata,
                                      attn_layernorm,
                                      self.attention,
                                      mlp_layernorm,
                                      self.feed_forward,
                                      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__internlm2__InternLM2Model__forward(
        self,
        input_ids: Optional[torch.Tensor],
        positions: torch.Tensor,
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
        intermediate_tensors: Optional[IntermediateTensors],
        inputs_embeds: Optional[torch.Tensor] = None,
    ) -> Union[torch.Tensor, IntermediateTensors]:
    '''
    =============================
    Modify by vllm_mlu
    =============================
    '''
    return decoder_model_forward_base_pp(input_ids, positions, kv_caches, attn_metadata, intermediate_tensors,
                                         self.layers, self.start_layer, self.end_layer,
                                         self.tok_embeddings,
                                         self.norm,
                                         inputs_embeds)
    '''
    ==================
    End of MLU Hijack
    ==================
    '''


MluHijackObject.apply_hijack(InternLM2Attention,
                             InternLM2Attention.forward,
                             vllm__module_executor__models__internlm2__InternLM2Attention__forward)
MluHijackObject.apply_hijack(InternLMDecoderLayer,
                             InternLMDecoderLayer.__init__,
                             vllm__module_executor__models__internlm2__InternLMDecoderLayer____init__)
MluHijackObject.apply_hijack(InternLM2ForCausalLM,
                             InternLM2ForCausalLM.load_weights,
                             vllm__module_executor__models__internlm2__InternLM2ForCausalLM__load_weights)
MluHijackObject.apply_hijack(InternLMDecoderLayer,
                             InternLMDecoderLayer.forward,
                             vllm__module_executor__models__internlm2__InternLMDecoderLayer__forward)
MluHijackObject.apply_hijack(InternLM2Model,
                             InternLM2Model.forward,
                             vllm__module_executor__models__internlm2__InternLM2Model__forward)
