"""
Inference-only MOE model.

Tensor Parallel evenly splits each expert's weight and distributes them to different ranks,
which means each rank holds partial weight of all experts.
While Expert Parallel evenly distributes some of the experts' full weight to different ranks,
which means each rank holds part of the experts' full weight.

As a result, each rank in the Tensor Parallel group receives all tokens' hidden states for all experts,
then computes using the partial weights, while for Expert Parallel, each rank only receives
part of tokens' hidden states for experts on this rank, then computes using the full weights.

When both Tensor Parallel and Expert Parallel are enabled, each rank handles
a portion of the expert weights matrices (as in EP mode) and these weights are further sliced
across ranks (as in TP mode). This hybrid approach aims to balance the workload more evenly across ranks,
enhancing efficiency and reducing the likelihood of bottlenecks associated with EP mode alone.
"""

from typing import Optional

import torch
from torch import nn

from vllm.distributed import (get_tensor_model_parallel_rank,
                              get_tensor_model_parallel_world_size,
                              get_tensor_model_parallel_group)
from examples.cambricon_custom_func.vllm.mlu_hijack.distributed.parallel_state import (
    get_moe_tensor_parallel_rank, get_moe_tensor_parallel_world_size, get_moe_tensor_parallel_group,
    get_moe_expert_parallel_rank, get_moe_expert_parallel_world_size, get_moe_expert_parallel_group)
from vllm.model_executor.layers.linear import ReplicatedLinear
from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
from vllm_mlu.model_executor.layers.feed_forward import FeedForward
from vllm_mlu.model_executor.layers.sparse_moe_mlp import SparseMoeMlp
from vllm_mlu.mlu_hijack_utils import MluHijackObject
from vllm_mlu._mlu_utils import get_device_major_capability


def vllm__mlu_hijack__model_executor__layers__feed_forward__SparseMoeMlp____init__(
        self,
        num_experts: int,
        top_k: int,
        hidden_size: int,
        intermediate_size: int,
        up_proj_name: str,
        is_gated: bool,
        down_proj_name: str,
        has_bias: bool,
        skip_bias_add: bool = False,
        renormalize:bool = False,
        hidden_act: str = "silu",
        params_dtype: Optional[torch.dtype] = None,
        quant_config: Optional[QuantizationConfig] = None,
        is_use_fused_moe: bool = False,
        expert_group: int = 1,
        topk_group: int = 1,
    ):
    super(SparseMoeMlp, self).__init__()
    self.tp_rank = get_tensor_model_parallel_rank()
    self.tp_size = get_tensor_model_parallel_world_size()
    self.tp_group = get_tensor_model_parallel_group()
    self.num_total_experts = num_experts
    self.top_k = top_k
    self.hidden_size = hidden_size
    self.intermediate_size = intermediate_size
    self.up_proj_name = up_proj_name
    self.is_gated = is_gated
    self.down_proj_name = down_proj_name
    self.has_bias = has_bias
    self.renormalize = renormalize
    self.hidden_act = hidden_act
    self.quant_config = quant_config
    self.is_use_fused_moe = is_use_fused_moe
    self.expert_group = expert_group
    self.topk_group = topk_group
    if get_device_major_capability() == 3:
        self.is_use_fused_moe = False

    if params_dtype is None:
        params_dtype = torch.get_default_dtype()
    self.params_dtype = params_dtype

    '''
    =============================
    Modify by vllm_mlu
    =============================
    @brief: add moe relative distribution
    '''
    self.moe_tp_size = get_moe_tensor_parallel_world_size()
    self.moe_tp_rank = get_moe_tensor_parallel_rank()
    self.moe_tp_group = get_moe_tensor_parallel_group()
    self.moe_ep_size = get_moe_expert_parallel_world_size()
    self.moe_ep_rank = get_moe_expert_parallel_rank()
    self.moe_ep_group = get_moe_expert_parallel_group()

    # NOTE: The bias for fc2 is only applied on tp_rank 0. If we added it on all nodes the allreduce() would
    # contain multiple copies of the bias. The bias on other node will be ignored, and may be set to nullptr
    self.skip_bias_add = True if self.moe_tp_rank > 0 else False

    assert self.num_total_experts >= self.moe_ep_size, (
        f"need num_total_experts:{self.num_total_experts} >= moe_ep_size:{self.moe_ep_size}")

    assert self.intermediate_size % self.moe_tp_size == 0, (
        f"need intermediate_size:{self.intermediate_size} % moe_tp_size:{self.moe_tp_size} == 0")

    self.num_experts_per_rank = (self.num_total_experts + self.moe_ep_size - 1) // self.moe_ep_size
    if self.moe_ep_rank + 1 == self.moe_ep_size and self.num_total_experts % self.moe_ep_size:
        self.num_experts_per_rank = self.num_total_experts % self.moe_ep_size

    self.start_expert_id = self.moe_ep_rank * ((self.num_total_experts + self.moe_ep_size - 1) // self.moe_ep_size)
    '''
    ==================
    End of MLU Hijack
    ==================
    '''
    self.end_expert_id = self.start_expert_id + self.num_experts_per_rank

    # Gate always runs at half / full precision for now.
    self.gate = ReplicatedLinear(self.hidden_size,
                                 self.num_total_experts,
                                 bias=False,
                                 params_dtype=self.params_dtype,
                                 quant_config=None)
    self.experts = nn.ModuleList([
        FeedForward(hidden_size=self.hidden_size,
                    intermediate_size=self.intermediate_size,
                    hidden_act=self.hidden_act,
                    up_proj_name=self.up_proj_name,
                    is_gated=self.is_gated,
                    down_proj_name=self.down_proj_name,
                    bias=self.has_bias,
                    quant_config=self.quant_config,
                    skip_bias_add=self.skip_bias_add,
                    reduce_results=False,
                    tp_group=self.moe_tp_group) for idx in range(self.num_experts_per_rank)
    ])

    self.init_pack_param()


MluHijackObject.apply_hijack(SparseMoeMlp,
                             SparseMoeMlp.__init__,
                             vllm__mlu_hijack__model_executor__layers__feed_forward__SparseMoeMlp____init__)
