import torch
from typing import Optional
from vllm_mlu.mlu_hijack_utils import MluHijackObject, set_is_gated
from vllm_mlu.model_executor.layers.feed_forward import FeedForward
from vllm.distributed.parallel_state import get_tp_group, get_tensor_model_parallel_group
from vllm.distributed.communication_op import tensor_model_parallel_all_reduce
from vllm.distributed import get_tensor_model_parallel_rank
from vllm import _mlu_ops as mlu_ops
from vllm.lora.layers import BaseLayerWithLoRA
from vllm_mlu._mlu_utils import *


def vllm_mlu__model_executor__layers__feed_forward__FeedForward__forward(
    self, 
    hidden_states, 
    residual: Optional[torch.Tensor] = None
):
    self.prepare_weight()
    up_proj = getattr(self, self.up_proj_name)
    down_proj = getattr(self, self.down_proj_name)
    residual_ = None if self.tp_rank > 0 else residual
    if (self.use_bt_ffn and not isinstance(up_proj, BaseLayerWithLoRA)
        and not isinstance(down_proj, BaseLayerWithLoRA)):
        # The matmul formula is the following:
        #   mul_out = alpha * (matmul(input, filter, transpose\_b=True) + bias) + beta * residual
        #   output = active(mul_out)
        # Notes: We cannot use the activation function in matmul because it does not support gated operation
        #  we might support its in tmo matmul in the future
        fc1 = mlu_ops.matmul(hidden_states.view(-1, self.hidden_size), up_proj.weight, up_proj.bias, 
                            None, 'none', self.alpha, self.beta)
        act_out = mlu_ops.active(fc1, self.hidden_act, self.is_gated)
        beta = 1.0 if residual_ is not None else 0.0
        '''
        =======================================
        Modify by custom vllm_mlu
        =======================================
        @brief: call parallel op and abandon original reduce if parallel_num is set
        '''
        is_parallel_enable = hasattr(self, 'parallel_num') and get_is_prompt()
        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)
            out_ = mlu_ops.matmul_allreduce(cncl_comm, act_out, down_proj.weight, None, residual_,
                                            self.alpha, beta, self.parallel_num)
        else:
            out_ = mlu_ops.matmul(act_out, down_proj.weight, None, residual_, 'none', self.alpha, beta)
        '''
        =======================================
        End of custom MLU Hijack
        =======================================
        '''
        # bias if existed need to add after second matmul according to the original design of vllm
        '''
        =============================
        Modify by custom vllm_mlu
        =============================
        @brief: when preload_size is set, call GroupCoordinator.all_reduce() directly and 
        use async_op to set all_reduce paralleled with preload 
        '''
        if self.reduce_results and self.tp_size > 1 and not is_parallel_enable:
            if hasattr(self, 'preload_size') and self.preload_size > 0 and not self.is_prompt:
                handle = get_tp_group().all_reduce(out_, async_op=True)
                _MB = 1 << 20
                mlu_ops.preload(self.preloaded_weights[0].data, self.preload_size * _MB)
                preloaded_weights_size = self.preloaded_weights[0].numel() * self.preloaded_weights[0].element_size()
                if preloaded_weights_size < (self.preload_size * _MB) and len(self.preloaded_weights) > 1:
                    mlu_ops.preload(self.preloaded_weights[1].data, (self.preload_size * _MB) - preloaded_weights_size)
                handle.wait()
                out = out_
            else:
                out = tensor_model_parallel_all_reduce(out_)
        else:
            out = out_
        '''
        =========================
        End of custom MLU Hijack
        =========================
        '''
        # do the bias add if needed
        if not self.skip_bias_add:
            out = out + down_proj.bias if down_proj.bias is not None else out
        else:
            return out, down_proj.bias
    else:
        fc1, bias = up_proj(hidden_states)
        if bias is not None:
            fc1 += bias
        fc1 = mlu_ops.active(fc1, self.hidden_act, self.is_gated)
        out, bias = down_proj(fc1, residual=residual_)
        if self.skip_bias_add:
            return out, bias
    return out


MluHijackObject.apply_hijack(FeedForward,
                             FeedForward.forward,
                             vllm_mlu__model_executor__layers__feed_forward__FeedForward__forward)
