"""Inference-only MOE model."""
import torch
from torch import nn
from typing import Optional
from vllm.distributed import tensor_model_parallel_all_reduce, get_tensor_model_parallel_rank
from vllm.distributed.parallel_state import get_tensor_model_parallel_group
from vllm import _mlu_ops as mlu_ops
from vllm_mlu._mlu_utils import *
from vllm_mlu.mlu_hijack_utils import MluHijackObject
from vllm_mlu.model_executor.layers.sparse_moe_mlp import SparseMoeMlp


def vllm_mlu__model_executor__layers__sparse_moe_mlp__SparseMoeMlp__forward(
    self, 
    hidden_states: torch.Tensor,
    residual: Optional[torch.Tensor] = None
) -> torch.Tensor:
    orig_hidden_states_shape = hidden_states.shape
    hidden_states = hidden_states.view(-1, self.hidden_size)
    # expert_logits: [num_tokens, self.num_experts_per_rank]
    expert_logits, _ = self.gate(hidden_states)
    final_hidden_states = self.forward_experts(hidden_states, expert_logits, residual)

    '''
    =====================================================
    Modify by Context Communication Computation Parallel
    =====================================================
    @brief: disbale reduce if parallel op used
    '''
    is_parallel_enable = hasattr(self, 'parallel_num') and get_is_prompt()
    if self.tp_size > 1 and not is_parallel_enable:
        final_hidden_states = tensor_model_parallel_all_reduce(
            final_hidden_states)
    '''
    =====================================================
    End of Context Communication Computation Parallel
    =====================================================
    '''

    output = final_hidden_states.view(orig_hidden_states_shape)
    return output


def vllm_mlu__model_executor__layers__sparse_moe_mlp__SparseMoeMlp__forward_experts(
    self, 
    hidden_states, 
    expert_logits,
    residual: Optional[torch.Tensor] = None
):
    residual_ = None if self.tp_rank > 0 else residual
    if self.is_use_fused_moe:
        self.pack_params()
        '''
        =====================================================
        Modify by Context Communication Computation Parallel
        =====================================================
        @brief: call fused_moe all_reduce
        '''
        is_parallel_enable = hasattr(self, 'parallel_num') and get_is_prompt()
        if is_parallel_enable:
            residual_ = residual
        params = [hidden_states, expert_logits, self.w13, self.w2, self.b13, self.b2, 
            residual_, self.a13_scale, self.a2_scale, self.w13_scale, self.w2_scale, 
            self.top_k, self.renormalize, self.is_gated, self.hidden_act, self.start_expert_id]
        if is_parallel_enable:
            rank = get_tensor_model_parallel_rank()
            pg = get_tensor_model_parallel_group().device_group
            cncl_comm = pg._get_backend(torch.device("mlu")).get_cncl_comm(rank)
            params.extend([self.parallel_num, cncl_comm])
        final_hidden_states = mlu_ops.fused_moe(*params)
        '''
        =====================================================
        End of Context Communication Computation Parallel
        =====================================================
        '''
    else:
        final_hidden_states = self.forward_experts_nofused(hidden_states, expert_logits)
        if residual_ is not None:
            final_hidden_states = final_hidden_states + residual_

    return final_hidden_states


MluHijackObject.apply_hijack(SparseMoeMlp,
                             SparseMoeMlp.forward,
                             vllm_mlu__model_executor__layers__sparse_moe_mlp__SparseMoeMlp__forward)
MluHijackObject.apply_hijack(SparseMoeMlp,
                             SparseMoeMlp.forward_experts,
                             vllm_mlu__model_executor__layers__sparse_moe_mlp__SparseMoeMlp__forward_experts)