import torch
from typing import Optional
from vllm import _mlu_ops as mlu_ops
from vllm_mlu._mlu_utils import *
from vllm_mlu.mlu_hijack_utils import MluHijackObject
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_mlu.model_executor.custom_model.custom import CustomMoeBlock


def vllm__module_executor__custom_model__CustomMoeBlock__forward(
    self, 
    hidden_states: torch.Tensor,
    residual: Optional[torch.Tensor] = None
) -> torch.Tensor:
    num_tokens, hidden_dim = hidden_states.shape
    hidden_states = hidden_states.view(-1, hidden_dim)
    shared_output = None
    if self.shared_expert is not None:
        shared_output = self.shared_expert(hidden_states)
        if self.shared_expert_gate is not None:
            shared_output = F.sigmoid(
                self.shared_expert_gate(hidden_states)) * shared_output

    # router_logits: (num_tokens, n_experts)
    router_logits, _ = self.gate(hidden_states)
    residual_ = None if self.rank > 0 else residual
    '''
    =====================================================
    Modify by Context Communication Computation Parallel
    =====================================================
    @brief: call fused_moe
    '''
    params = [hidden_states, router_logits, self.w1, self.w2, None, None, 
        residual_, self.input_smooth, self.act_smooth, self.w1_scale, self.w2_scale, 
        self.top_k, self.config.norm_topk_prob, self.config.is_gated, self.config.hidden_act, 0]
    if hasattr(self, 'parallel_num') and get_is_prompt():
        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
    =====================================================
    '''

    if shared_output is not None:
        final_hidden_states = final_hidden_states + shared_output
        
    reduce_results = (self.config.use_parallel_residual == False)
    if reduce_results:
        final_hidden_states = tensor_model_parallel_all_reduce(
            final_hidden_states)

    return final_hidden_states.view(num_tokens, hidden_dim)


MluHijackObject.apply_hijack(CustomMoeBlock,
                             CustomMoeBlock.forward,
                             vllm__module_executor__custom_model__CustomMoeBlock__forward)
