import torch
import torch.nn.functional as F
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 get_tensor_model_parallel_rank, tensor_model_parallel_all_reduce
from vllm.distributed.parallel_state import get_tensor_model_parallel_group
from vllm.attention import AttentionMetadata
from vllm_mlu.model_executor.models.qwen2_moe import Qwen2MoeSparseMoeBlock


def vllm_mlu__model_executor__models__qwen2_moe__Qwen2MoeSparseMoeBlock__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:
            gate_output = self.shared_expert_gate(hidden_states)
            shared_output = F.sigmoid(gate_output[0]) * shared_output

    # router_logits: (num_tokens, n_experts)
    router_logits, _ = self.gate(hidden_states)
    final_hidden_states = self.forward_experts(hidden_states, router_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:
        if is_parallel_enable:
            shared_output = tensor_model_parallel_all_reduce(shared_output)
            if shared_output is not None:
                final_hidden_states = final_hidden_states + shared_output
        else:
            if shared_output is not None:
                final_hidden_states = final_hidden_states + shared_output
            final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states)
    '''
    =====================================================
    End of Context Communication Computation Parallel
    =====================================================
    '''

    return final_hidden_states.view(num_tokens, hidden_dim)


MluHijackObject.apply_hijack(Qwen2MoeSparseMoeBlock,
                             Qwen2MoeSparseMoeBlock.forward,
                             vllm_mlu__model_executor__models__qwen2_moe__Qwen2MoeSparseMoeBlock__forward)