import torch
import torch.nn.functional as F
from typing import Optional
from vllm.config import CacheConfig
from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
from vllm_mlu._mlu_utils import *
from vllm_mlu.model_executor.layers.feed_forward import FeedForward
from vllm_mlu.mlu_hijack_utils import MluHijackObject
from vllm.distributed import tensor_model_parallel_all_reduce
from vllm_mlu.transformers_utils.configs import CustomConfig
from vllm_mlu.model_executor.custom_model.custom import CustomDecoderLayer, CustomAttention, _NORM_DICT
from vllm_mlu.model_executor.layers.sparse_moe_mlp import SparseMoeMlp
from vllm.model_executor.layers.linear import ReplicatedLinear
from vllm_mlu.model_executor.models.layer_utils import (
    decoder_layer_forward_base, is_per_tensor_smoothquant,
    is_per_token_smoothquant, quant_fusion_with_rmsnorm,
    quant_fusion_with_layernorm)


class CustomMoeBlock(SparseMoeMlp):

    def __init__(
        self,
        config: CustomConfig,
        quant_config: Optional[QuantizationConfig] = None,
    ):
        super().__init__(num_experts=config.num_experts,
                         top_k=config.num_experts_per_tok,
                         hidden_size=config.hidden_size,
                         intermediate_size=config.moe_intermediate_size,
                         up_proj_name="gate_up_proj",
                         is_gated=config.is_gated,
                         down_proj_name="down_proj",
                         has_bias=config.mlp_bias,
                         skip_bias_add=False,
                         renormalize=config.norm_topk_prob,
                         hidden_act=config.hidden_act,
                         params_dtype=None,
                         quant_config=quant_config,
                         is_use_fused_moe=True)

        self.config = config
        self.rank = self.tp_rank
        self.shared_expert = None
        self.shared_expert_gate = None
        if config.shared_expert_intermediate_size > 0:
            self.shared_expert = FeedForward(hidden_size=config.hidden_size,
                                             intermediate_size=config.shared_expert_intermediate_size,
                                             hidden_act=config.hidden_act,
                                             up_proj_name='gate_up_proj',
                                             is_gated=config.is_gated,
                                             down_proj_name='down_proj',
                                             bias=config.mlp_bias,
                                             quant_config=quant_config,
                                             reduce_results=False)
            self.shared_expert_gate = ReplicatedLinear(config.hidden_size,
                                                       1,
                                                       bias=False,
                                                       params_dtype=self.params_dtype,
                                                       quant_config=None)


    def 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)
        residual_ = None if self.rank > 0 else residual
        '''
        =============================
        Modify by vllm_mlu
        =============================
        @brief: modify bt_ops.fused_moe to forward_experts
        '''
        final_hidden_states = self.forward_experts(hidden_states, router_logits, residual)
        '''
        ==================
        End of MLU Hijack
        ==================
        '''
        if shared_output is not None:
            final_hidden_states = final_hidden_states + shared_output
        '''
        =============================
        Modify by vllm_mlu
        =============================
        @brief: add comment to explain use_parallel_residual usage
        '''
        # use_parallel_residual = True: x = x + attn(ln1(x)) + mlp(ln2(x))
        # use_parallel_residual = False:
        #   if apply_residual_connection_post_layernorm:
        #       x_attn = ln1(x) + attn(ln1(x))
        #       x_mlp = ln2(x_attn) + mlp(ln2(x_attn))
        #   else:
        #       x_attn = x + attn(ln1(x))
        #       x_mlp = x_attn + mlp(ln2(x_attn))
        # When use_parallel_residual = True, x is shared between attn and mlp, so we only need to
        # reduce after x + attn(ln1(x)) + mlp(ln2(x)) and don't need reduce here
        # But when use_parallel_residual = False, mlp layer uses attn layer's output, so need reduce
        # when mlp is finished.
        '''
        ==================
        End of MLU Hijack
        ==================
        '''
        reduce_results = (self.config.use_parallel_residual == False)
        if reduce_results and self.tp_size > 1:
            final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states)

        return final_hidden_states.view(num_tokens, hidden_dim)


def vllm__mlu_hijack__model_executor__custom_model__custom__CustomDecoderLayer____init__(
        self,
        config: CustomConfig,
        cache_config: Optional[CacheConfig] = None,
        quant_config: Optional[QuantizationConfig] = None,
    ) -> None:
    super(CustomDecoderLayer, self).__init__()
    self.config = config
    self.self_attn = CustomAttention(
        config=config,
        cache_config=cache_config,
        quant_config=quant_config,
    )

    mlp_bias = getattr(config, "mlp_bias", False) or getattr(config, "bias", False)
    is_gated = getattr(config, "is_gated", False)

    if config.num_experts is not None:
        '''
            =============================
            Modify by vllm_mlu
            =============================
            @brief: nothing changed, only use the CustomMoeBlock class in this file
            '''
        self.mlp = CustomMoeBlock(config=config,
                                quant_config=quant_config)
        '''
            ==================
            End of MLU Hijack
            ==================
            '''
    else:
        self.mlp = FeedForward(hidden_size=config.hidden_size,
                               intermediate_size=config.intermediate_size,
                               hidden_act=self.config.hidden_act,
                               up_proj_name='up_proj',
                               is_gated=is_gated,
                               down_proj_name='down_proj',
                               bias=mlp_bias,
                               quant_config=quant_config,
                               skip_bias_add=(self.config.use_parallel_residual and mlp_bias),
                               reduce_results = (self.config.use_parallel_residual == False))

    self.input_layernorm = _NORM_DICT[self.config.norm_type](config.hidden_size, eps=config.norm_eps)
    self.post_attention_layernorm = _NORM_DICT[self.config.norm_type](config.hidden_size, eps=config.norm_eps)

    # perf per-tensor sq cases by fusing quantization in layernorm
    self.is_per_tesnor_sq_perf_cases = (is_per_tensor_smoothquant(quant_config) and
                                        not self.config.apply_residual_connection_post_layernorm)
    self.is_per_token_sq_perf_cases = (is_per_token_smoothquant(quant_config) and
                                        not self.config.apply_residual_connection_post_layernorm)
    if self.is_per_tesnor_sq_perf_cases or self.is_per_token_sq_perf_cases:
        self.self_attn.qkv_proj.quant_method.skip_quant_input = True
        self.quant_fusion_attn_layernorm = None
        self.is_moe = config.num_experts is not None
        self.use_rmsnorm = self.config.norm_type == "rmsnorm"
        if not self.is_moe:
            self.mlp.up_proj.quant_method.skip_quant_input = True
            self.quant_fusion_mlp_layernorm = None


MluHijackObject.apply_hijack(CustomDecoderLayer,
                             CustomDecoderLayer.__init__,
                             vllm__mlu_hijack__model_executor__custom_model__custom__CustomDecoderLayer____init__)
