import torch
import re
import vllm
from torch import nn
from typing import List, Optional, Tuple, Iterable
from vllm_mlu._mlu_utils import *
from vllm_mlu.mlu_hijack_utils import MluHijackObject
from vllm.model_executor.models.mixtral import MixtralForCausalLM
from vllm.model_executor.model_loader.weight_utils import default_weight_loader, maybe_remap_kv_scale_name
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.model_executor.models.utils import is_pp_missing_parameter


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()
    # 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_local_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"),
        ("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)


MluHijackObject.apply_hijack(MixtralForCausalLM,
                             MixtralForCausalLM.load_weights,
                             vllm__module_executor__models__mixtral__MixtralForCausalLM__load_weights)
