import torch
import re
from typing import Optional, Iterable, Tuple
from vllm_mlu._mlu_utils import *
from vllm_mlu.mlu_hijack_utils import MluHijackObject
from vllm.model_executor.models.qwen2_moe import  Qwen2MoeForCausalLM
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
from examples.cambricon_custom_func.vllm.mlu_hijack.distributed.parallel_state import (
    get_moe_expert_parallel_rank, get_moe_expert_parallel_world_size)
from vllm_mlu.model_executor.layers.sparse_moe_mlp import SparseMoeMlp
from vllm.utils import print_warning_once
from vllm.model_executor.models.utils import is_pp_missing_parameter


def vllm__module_executor__models__qwen2moe__Qwen2MoeForCausalLM__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.num_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)
        ("qkv_proj", "q_proj", "q"),
        ("qkv_proj", "k_proj", "k"),
        ("qkv_proj", "v_proj", "v"),
        ("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:
            if weight_name not in name:
                continue
            '''
            =============================
            Modify by vllm_mlu
            =============================
            @brief: delete if "mlp.experts" in name: continue condition
            '''
            '''
            ==================
            End of MLU Hijack
            ==================
            '''
            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
            # Skip layers on other devices.
            if is_pp_missing_parameter(name, self):
                continue
            '''
            =============================
            Modify by vllm_mlu
            =============================
            @brief: add expert skiped condition and delete useless if name not in params_dict: continue condition
            '''
            # Skip experts that are not assigned to this worker.
            if (("mlp.experts." in name or "mlp.shared_expert." in name or "mlp.shared_expert_gate." 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:
            '''
            =============================
            Modify by vllm_mlu
            =============================
            @brief: delete for mapping in expert_params_mapping condition
            '''
            '''
            ==================
            End of MLU Hijack
            ==================
            '''
            # Skip loading extra bias for GPTQ models.
            if ((name.endswith(".bias") or name.endswith("_bias"))
                            and name not in params_dict):
                continue
            # Skip layers on other devices.
            if is_pp_missing_parameter(name, self):
                continue
            # Remapping the name of FP8 kv-scale.
            if name.endswith("kv_scale"):
                remapped_kv_scale_name = name.replace(
                    ".kv_scale", ".attn.kv_scale")
                if remapped_kv_scale_name not in params_dict:
                    print_warning_once(
                        "Found kv scale in the checkpoint "
                        f"(e.g. {name}), but not found the expected "
                        f"name in the model "
                        f"(e.g. {remapped_kv_scale_name}). "
                        "kv-scale is not loaded.")
                    continue
                else:
                    name = remapped_kv_scale_name
            '''
            =============================
            Modify by vllm_mlu
            =============================
            @brief: add expert skiped condition
            '''
            # Skip experts that are not assigned to this worker.
            if (("mlp.experts." in name or "mlp.shared_expert." in name or "mlp.shared_expert_gate." 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)


MluHijackObject.apply_hijack(Qwen2MoeForCausalLM,
                             Qwen2MoeForCausalLM.load_weights,
                             vllm__module_executor__models__qwen2moe__Qwen2MoeForCausalLM__load_weights)
