
import re
import torch
from torch import nn
from typing import Any, Dict, Iterable, List, Optional, Tuple
from transformers import PretrainedConfig
from vllm.model_executor.layers.linear import (ColumnParallelLinear,
                                               MergedColumnParallelLinear,
                                               ReplicatedLinear,
                                               RowParallelLinear)
from vllm.model_executor.layers.quantization.base_config import (
    QuantizationConfig)
from vllm.model_executor.model_loader.weight_utils import default_weight_loader

from vllm_mlu.model_executor.layers.feed_forward import FeedForward
from vllm_mlu.mlu_hijack_utils import MluHijackObject
from vllm_mlu.model_executor.layers.sparse_moe_mlp import SparseMoeMlp
from vllm.model_executor.models.utils import is_pp_missing_parameter

from vllm.model_executor.models.deepseek_v2 import DeepseekV2ForCausalLM
from vllm_mlu.model_executor.models.deepseek_v2  import DeepseekV2MoE
from examples.cambricon_custom_func.vllm.mlu_hijack.distributed.parallel_state import (
    get_moe_expert_parallel_rank, get_moe_expert_parallel_world_size)


def vllm_mlu__model_executor__models__deepseek_v2__DeepseekV2MoE____init__(
        self,
        config: PretrainedConfig,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
    ):
    super(DeepseekV2MoE, self).__init__(num_experts=config.n_routed_experts,
                     top_k=config.num_experts_per_tok,
                     hidden_size=config.hidden_size,
                     intermediate_size=config.moe_intermediate_size,
                     up_proj_name="gate_up_proj",
                     is_gated=True,
                     down_proj_name="down_proj",
                     has_bias=False,
                     skip_bias_add=False,
                     renormalize=config.norm_topk_prob,
                     hidden_act=config.hidden_act,
                     params_dtype=None,
                     quant_config=quant_config,
                     is_use_fused_moe=True,
                     expert_group=config.n_group,
                     topk_group=config.topk_group)
    self.config = config
    self.routed_scaling_factor = config.routed_scaling_factor
    self.n_shared_experts = config.n_shared_experts
    self.routed_scaling_factor = config.routed_scaling_factor
    if self.moe_tp_size > config.n_routed_experts:
        raise ValueError(
            f"Moe Tensor parallel size {self.moe_tp_size} is greater than "
            f"the number of experts {config.n_routed_experts}.")

    if config.hidden_act != "silu":
        raise ValueError(f"Unsupported activation: {config.hidden_act}. "
                         "Only silu is supported for now.")

    self.gate = ReplicatedLinear(config.hidden_size,
                                 config.n_routed_experts,
                                 bias=False,
                                 quant_config=None,
                                 prefix=f"{prefix}.gate")
    if config.n_shared_experts is not None:
        intermediate_size = (config.moe_intermediate_size *
                             config.n_shared_experts)
        '''
        =============================
        Modify by vllm_mlu
        =============================
        @brief: replace MLP with FeedForward.
        '''
        self.shared_experts = FeedForward(hidden_size=config.hidden_size,
                                         intermediate_size=intermediate_size,
                                         hidden_act=config.hidden_act,
                                         up_proj_name='gate_up_proj',
                                         is_gated=True,
                                         down_proj_name='down_proj',
                                         bias=False,
                                         quant_config=quant_config,
                                         reduce_results=False)
        '''
        ==================
        End of MLU Hijack
        ==================
        '''


def vllm__module_executor__models__deepseek_v2__DeepseekV2ForCausalLM__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()

    # expert parallel modification start
    moe_ep_rank = get_moe_expert_parallel_rank()
    moe_ep_size = get_moe_expert_parallel_world_size()
    num_total_experts = self.config.n_routed_experts
    start_expert_id = moe_ep_rank * ((num_total_experts + moe_ep_size - 1) // moe_ep_size)
    # expert parallel modification end
    '''
    ==================
    End of MLU Hijack
    ==================
    '''

    stacked_params_mapping = [
        # (param_name, shard_name, shard_id)
        ("gate_up_proj", "gate_proj", 0),
        ("gate_up_proj", "up_proj", 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 "mlp.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:
            # Skip non-stacked layers and experts (experts handled below).
            if weight_name not in name:
                continue
            # We have mlp.experts[0].gate_proj in the checkpoint.
            # Since we handle the experts below in expert_params_mapping,
            # we need to skip here BEFORE we update the name, otherwise
            # name will be updated to mlp.experts[0].gate_up_proj, which
            # will then be updated below in expert_params_mapping
            # for mlp.experts[0].gate_gate_up_proj, which breaks load.
            '''
            =============================
            Modify by vllm_mlu
            =============================
            @brief: add expert skiped condition and delete useless if name not in params_dict: continue condition
            '''
            name = name.replace(weight_name, param_name)
            if (("mlp.experts." in name or "mlp.shared_experts." in name or "mlp.shared_expert_gate." in name)
                    and name not in params_dict):
                continue
            '''
            ==================
            End of MLU Hijack
            ==================
            '''
            # 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:
            '''
            =============================
            Modify by vllm_mlu
            =============================
            @brief: add expert skiped condition
            '''
            # Skip loading extra bias for GPTQ models.
            if name.endswith(".bias") and name not in params_dict:
                continue

            if (("mlp.experts." in name or "mlp.shared_experts." in name or "mlp.shared_expert_gate." in name)
                    and name not in params_dict):
                continue
            if is_pp_missing_parameter(name, self):
                continue

            param = params_dict[name]
            weight_loader = getattr(param, "weight_loader",
                                    default_weight_loader)
            weight_loader(param, loaded_weight)
            '''
            ==================
            End of MLU Hijack
            ==================
            '''


MluHijackObject.apply_hijack(DeepseekV2MoE,
                             DeepseekV2MoE.__init__,
                             vllm_mlu__model_executor__models__deepseek_v2__DeepseekV2MoE____init__)
MluHijackObject.apply_hijack(DeepseekV2ForCausalLM,
                             DeepseekV2ForCausalLM.load_weights,
                             vllm__module_executor__models__deepseek_v2__DeepseekV2ForCausalLM__load_weights)
