from typing import Optional, Tuple
import torch
import vllm
from vllm_mlu.mlu_hijack_utils import MluHijackObject
from vllm_mlu.model_executor.layers.rotary_embedding import MLURotaryEmbedding
from examples.cambricon_custom_func.vllm.mlu_hijack.distributed.parallel_state import (
    get_context_model_parallel_world_size)

def vllm__module_executor__layers__rotary_embedding__MLURotaryEmbedding__forward_mlu_wraper(
    self,
    positions: torch.Tensor,
    x: torch.Tensor,
    offsets: Optional[torch.Tensor] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
    from vllm import _mlu_ops as mlu_ops

    # ops.rotary_embedding()/batched_rotary_embedding()
    # are in-place operations that update the query and key tensors.
    if offsets is not None:
        raise ValueError(f"tmo.apply_rotary not support offsets yet.")
    else:
        if MLURotaryEmbedding.set_cos_sin == False:
            MLURotaryEmbedding.cos_, MLURotaryEmbedding.sin_ = self._get_cos_sin()
            MLURotaryEmbedding.set_cos_sin = True
        interleaved = True
        if self.is_neox_style:
            interleaved = False
        if MLURotaryEmbedding.is_chunked or not MLURotaryEmbedding.is_prompt:
            position_ids = positions
            discrete = True
        else :
            position_ids = None
            discrete = False
        '''
        ==========================
        Modify by Context Parallel
        ==========================
        @brief: context parallel need discrete = True
        '''
        position_ids = None if (MLURotaryEmbedding.is_prompt and get_context_model_parallel_world_size == 1) else positions
        discrete = False if (MLURotaryEmbedding.is_prompt and get_context_model_parallel_world_size == 1) else True
        '''
        =======================
        End of Context Parallel
        =======================
        '''
        x = mlu_ops.rotary_embedding(x,
                                    MLURotaryEmbedding.sin_,
                                    MLURotaryEmbedding.cos_,
                                    position_ids,
                                    MLURotaryEmbedding.cu_seq_lens,
                                    interleaved,
                                    discrete,
                                    False,
                                    MLURotaryEmbedding.max_seq_len)

    return x


MluHijackObject.apply_hijack(MLURotaryEmbedding,
                             MLURotaryEmbedding.forward_mlu,
                             vllm__module_executor__layers__rotary_embedding__MLURotaryEmbedding__forward_mlu_wraper)
