from typing import Optional, Any
from vllm.logger import init_logger
from vllm.model_executor.layers.quantization.base_config import (
    QuantizationConfig)
from vllm.model_executor.layers.linear import (
    MergedColumnParallelLinear,
    ColumnParallelLinear,
    RowParallelLinear
)
from vllm_mlu.mlu_hijack_utils import set_is_gated, MluHijackObject
from vllm_mlu.model_executor.layers.feed_forward import FeedForward
from ....mlu_hijack.distributed.parallel_state import (get_parallel_rank_with_group, get_parallel_world_size_with_group)


logger = init_logger(__name__)


def vllm__mlu_hijack__model_executor__layers__feed_forward__FeedForward____init__(
        self,
        hidden_size: int,
        intermediate_size: int,
        hidden_act: str,
        up_proj_name: str,
        is_gated: bool,
        down_proj_name: str,
        bias: bool,
        quant_config: Optional[QuantizationConfig] = None,
        skip_bias_add: bool = False,
        reduce_results: bool = True,
        prefix: str = "",
        tp_group: Any = None,
    ):
    super(FeedForward, self).__init__()
    self.hidden_size = hidden_size
    self.hidden_act = hidden_act
    self.is_gated = is_gated
    self.bias = bias
    self.up_proj_name = up_proj_name
    self.down_proj_name = down_proj_name
    self.quant_config = quant_config
    self.is_initialized = False
    self.skip_bias_add = skip_bias_add
    self.reduce_results = reduce_results
    self.use_bt_ffn = True if quant_config is None else False
    set_is_gated(self.is_gated)
    self.tp_size = get_parallel_world_size_with_group(tp_group)
    self.tp_rank = get_parallel_rank_with_group(tp_group)

    '''
        =============================
        Modify by vllm_mlu
        =============================
        @brief: add tp_group parameter at the end of each linear class
        '''
    self.tp_group = tp_group
    # up_proj with gate or not
    if self.is_gated:
        up_proj = MergedColumnParallelLinear(hidden_size,
                                             [intermediate_size] * 2,
                                             bias=bias,
                                             quant_config=quant_config,
                                             prefix=f"{prefix}.{up_proj_name}",
                                             tp_group=tp_group)
    else:
        up_proj = ColumnParallelLinear(hidden_size,
                                       intermediate_size,
                                       bias=bias,
                                       skip_bias_add=skip_bias_add,
                                       quant_config=quant_config,
                                       prefix=f"{prefix}.{up_proj_name}",
                                       tp_group=tp_group)
    self.register_module(up_proj_name, up_proj)

    # down_proj
    down_proj = RowParallelLinear(intermediate_size,
                                  hidden_size,
                                  bias=bias,
                                  skip_bias_add=skip_bias_add,
                                  reduce_results=reduce_results,
                                  quant_config=quant_config,
                                  prefix=f"{prefix}.{down_proj_name}",
                                  tp_group=tp_group)
    '''
        ==================
        End of MLU Hijack
        ==================
        '''
    self.register_module(down_proj_name, down_proj)


MluHijackObject.apply_hijack(FeedForward,
                             FeedForward.__init__,
                             vllm__mlu_hijack__model_executor__layers__feed_forward__FeedForward____init__)
