import torch
import re

from typing import List, Optional, Tuple, Union, Iterable
from vllm.attention import AttentionMetadata
from vllm.config import CacheConfig
from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
from vllm_mlu.mlu_hijack_utils import MluHijackObject
from vllm.model_executor.models.mixtral import (MixtralAttention, MixtralDecoderLayer,
                                                MixtralForCausalLM, MixtralModel)
from vllm_mlu.mlu_hijack_utils import set_is_gated
from transformers import MixtralConfig
from vllm_mlu.model_executor.layers.sparse_moe_mlp import SparseMoeMlp
from vllm.model_executor.layers.layernorm import RMSNorm
from vllm.model_executor.model_loader.weight_utils import (default_weight_loader,
                                                           maybe_remap_kv_scale_name)
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)


def vllm__module_executor__models__mixtral__MixtralAttention__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.qkv_proj(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(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.o_proj(attn_output, residual)
    '''
    ==================
    End of MLU Hijack
    ==================
    '''
    return output


def vllm__module_executor__models__mixtral__MixtralDecoderLayer____init__(
        self,
        config: MixtralConfig,
        cache_config: Optional[CacheConfig] = None,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
    ) -> None:
    super(MixtralDecoderLayer, self).__init__()
    self.hidden_size = config.hidden_size
    # Requires transformers > 4.32.0
    rope_theta = getattr(config, "rope_theta", 10000)
    self.self_attn = MixtralAttention(
        hidden_size=self.hidden_size,
        num_heads=config.num_attention_heads,
        max_position=config.max_position_embeddings,
        num_kv_heads=config.num_key_value_heads,
        rope_theta=rope_theta,
        cache_config=cache_config,
        quant_config=quant_config,
        prefix=f"{prefix}.self_attn")
    '''
    =============================
    Modify by vllm_mlu
    =============================
    @brief: replace MixtralMoE to SparseMoeMlp
    '''
    self.block_sparse_moe = SparseMoeMlp(num_experts=config.num_local_experts,
                                         top_k=config.num_experts_per_tok,
                                         hidden_size=config.hidden_size,
                                         intermediate_size=config.intermediate_size,
                                         up_proj_name="w13",
                                         is_gated=True,
                                         down_proj_name="w2",
                                         has_bias=False,
                                         skip_bias_add=False,
                                         renormalize=True,
                                         hidden_act=config.hidden_act,
                                         params_dtype=None,
                                         quant_config=quant_config,
                                         is_use_fused_moe=True)
    '''
    ==================
    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. MoE gate linear always runs
        in half/full precision for now, so we only do quant fusion in attn block.
    '''
    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.qkv_proj.quant_method.skip_quant_input = True
        self.quant_fusion_attn_layernorm = None
    '''
    ==================
    End of MLU Hijack
    ==================
    '''


def vllm__module_executor__models__mixtral__MixtralForCausalLM__load_weights(
        self,
        weights: Iterable[Tuple[str, torch.Tensor]]):
    '''
    =============================
    Modify by vllm_mlu
    =============================
    @brief: pack params and cal start expert id
    '''
    for name, m in self.model.named_modules():
        if isinstance(m, SparseMoeMlp):
            m.pack_params()

    start_expert_id = 0
    '''
    ==================
    End of MLU Hijack
    ==================
    '''
    stacked_params_mapping = [
        # (param_name, shard_name, shard_id)
        ("qkv_proj", "q_proj", "q"),
        ("qkv_proj", "k_proj", "k"),
        ("qkv_proj", "v_proj", "v"),
        ("w13", "w1", 0),
        ("w13", "w3", 1),
        ]

    '''
    =============================
    Modify by vllm_mlu
    =============================
    @brief: delete expert_params_mapping for no useless
    '''
    '''
    ==================
    End of MLU Hijack
    ==================
    '''

    params_dict = dict(self.named_parameters())
    for name, loaded_weight in weights:
        if "rotary_emb.inv_freq" in name:
            continue
        '''
        =============================
        Modify by vllm_mlu
        =============================
        @brief: replace expert_id in weight to named_expert_id in params_dict
        '''
        if start_expert_id > 0 and "block_sparse_moe.experts." in name:
            expert_str = re.search(r'experts\.\d+', name).group(0)
            expert_id=int(expert_str.split(".")[1])
            named_expert_id = expert_id - start_expert_id
            old_expert_name = f"experts.{expert_id}"
            new_expert_name = f"experts.{named_expert_id}"
            name = name.replace(old_expert_name, new_expert_name)
        '''
        ==================
        End of MLU Hijack
        ==================
        '''
        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") or name.endswith("_bias"))
                        and name not in params_dict):
                continue
            if is_pp_missing_parameter(name, self):
                continue
            '''
            =============================
            Modify by vllm_mlu
            =============================
            @brief: add expert skiped condition
            '''
            # Skip experts that are not assigned to this worker.
            if (("block_sparse_moe.experts." in name) and (name not in params_dict)):
                continue
            '''
            ==================
            End of MLU Hijack
            ==================
            '''
            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") or name.endswith("_bias"))
                            and name not in params_dict):
                continue
            if is_pp_missing_parameter(name, self):
                continue
            # Remapping the name of FP8 kv-scale.
            name = maybe_remap_kv_scale_name(name, params_dict)
            if name is None:
                continue
            '''
            =============================
            Modify by vllm_mlu
            =============================
            @brief: add expert skiped condition
            '''
            # Skip experts that are not assigned to this worker.
            if (("block_sparse_moe.experts." in name) and (name not in params_dict)):
                continue
            '''
            ==================
            End of MLU Hijack
            ==================
            '''
            param = params_dict[name]
            weight_loader = getattr(param, "weight_loader",
                                    default_weight_loader)
            weight_loader(param, loaded_weight)


def vllm__module_executor__models__mixtral__MixtralDecoderLayer__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
    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.qkv_proj.smooth
            else:
                attn_quant_scale = self.self_attn.qkv_proj.scale_to_int
            self.quant_fusion_attn_layernorm = quant_fusion_with_rmsnorm(
                self.input_layernorm, attn_quant_scale,
                dynamic_quant=self.is_per_token_sq_perf_cases)
        attn_layernorm = self.quant_fusion_attn_layernorm

    return decoder_layer_forward_base(positions, hidden_states, kv_cache, attn_metadata,
                                      attn_layernorm,
                                      self.self_attn,
                                      self.post_attention_layernorm,
                                      self.block_sparse_moe,
                                      input_norm_fuse_en=self.is_per_token_sq_perf_cases)
    '''
    ==================
    End of MLU Hijack
    ==================
    '''


def vllm__module_executor__models__mixtral__MixtralModel__forward(
        self,
        input_ids: Optional[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, positions, kv_caches, attn_metadata, intermediate_tensors,
                                         self.layers, self.start_layer, self.end_layer,
                                         self.embed_tokens,
                                         self.norm)
    '''
    ==================
    End of MLU Hijack
    ==================
    '''


set_is_gated(True)
MluHijackObject.apply_hijack(MixtralAttention,
                             MixtralAttention.forward,
                             vllm__module_executor__models__mixtral__MixtralAttention__forward)
MluHijackObject.apply_hijack(MixtralDecoderLayer,
                             MixtralDecoderLayer.__init__,
                             vllm__module_executor__models__mixtral__MixtralDecoderLayer____init__)
MluHijackObject.apply_hijack(MixtralForCausalLM,
                             MixtralForCausalLM.load_weights,
                             vllm__module_executor__models__mixtral__MixtralForCausalLM__load_weights)
MluHijackObject.apply_hijack(MixtralDecoderLayer,
                             MixtralDecoderLayer.forward,
                             vllm__module_executor__models__mixtral__MixtralDecoderLayer__forward)
MluHijackObject.apply_hijack(MixtralModel,
                             MixtralModel.forward,
                             vllm__module_executor__models__mixtral__MixtralModel__forward)

