"""Inference-only MOE model."""
from typing import Optional

import torch
from torch import nn

from vllm.distributed import (get_tensor_model_parallel_rank,
                              get_tensor_model_parallel_world_size,
                              get_tensor_model_parallel_group,
                              tensor_model_parallel_all_reduce)
from vllm.model_executor.layers.linear import ReplicatedLinear
from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
from vllm.model_executor.layers.quantization.fp8 import Fp8Config
from vllm_mlu.model_executor.layers.feed_forward import FeedForward
from vllm import _mlu_ops as mlu_ops
from vllm_mlu.model_executor.layers.quantization.weightonly import WeightOnlyConfig
from vllm_mlu.model_executor.layers.quantization.smoothquant import SmoothQuantConfig
from vllm_mlu._mlu_utils import get_device_major_capability
from vllm.platforms import current_platform

class SparseMoeMlp(nn.Module):
    """
    Tensor Parallel evenly splits each expert's weight and distributes them to different ranks,
    which means each rank holds partial weight of all experts.
    While Expert Parallel evenly distributes some of the experts' full weight to different ranks,
    which means each rank holds part of the experts' full weight.

    As a result, each rank in the Tensor Parallel group receives all tokens' hidden states for all experts,
    then computes using the partial weights, while for Expert Parallel, each rank only receives
    part of tokens' hidden states for experts on this rank, then computes using the full weights.

    When both Tensor Parallel and Expert Parallel are enabled, each rank handles
    a portion of the expert weights matrices (as in EP mode) and these weights are further sliced
    across ranks (as in TP mode). This hybrid approach aims to balance the workload more evenly across ranks,
    enhancing efficiency and reducing the likelihood of bottlenecks associated with EP mode alone.
    """

    def __init__(
        self,
        num_experts: int,
        top_k: int,
        hidden_size: int,
        intermediate_size: int,
        up_proj_name: str,
        is_gated: bool,
        down_proj_name: str,
        has_bias: bool,
        skip_bias_add: bool = False,
        renormalize:bool = False,
        hidden_act: str = "silu",
        params_dtype: Optional[torch.dtype] = None,
        quant_config: Optional[QuantizationConfig] = None,
        is_use_fused_moe: bool = False,
        expert_group: Optional[int] = 1,
        topk_group: Optional[int] = 1,
    ):
        super().__init__()
        self.tp_rank = get_tensor_model_parallel_rank()
        self.tp_size = get_tensor_model_parallel_world_size()
        self.tp_group = get_tensor_model_parallel_group()
        self.num_total_experts = num_experts
        self.top_k = top_k
        self.hidden_size = hidden_size
        self.intermediate_size = intermediate_size
        self.up_proj_name = up_proj_name
        self.is_gated = is_gated
        self.down_proj_name = down_proj_name
        self.has_bias = has_bias
        self.renormalize = renormalize
        self.hidden_act = hidden_act
        self.quant_config = quant_config
        self.is_use_fused_moe = is_use_fused_moe
        self.expert_group = expert_group
        self.topk_group = topk_group
        if get_device_major_capability() == 3:
            self.is_use_fused_moe = False

        if params_dtype is None:
            params_dtype = torch.get_default_dtype()
        self.params_dtype = params_dtype

        # NOTE: The bias for fc2 is only applied on tp_rank 0. If we added it on all nodes the allreduce() would
        # contain multiple copies of the bias. The bias on other node will be ignored, and may be set to nullptr
        self.skip_bias_add = True if self.tp_rank > 0 else False

        assert self.intermediate_size % self.tp_size == 0, (
            f"need intermediate_size:{self.intermediate_size} % tp_size:{self.tp_size} == 0")

        self.num_experts_per_rank = self.num_total_experts

        self.start_expert_id = 0
        self.end_expert_id = self.start_expert_id + self.num_experts_per_rank

        # Gate always runs at half / full precision for now.
        self.gate = ReplicatedLinear(self.hidden_size,
                                     self.num_total_experts,
                                     bias=False,
                                     params_dtype=self.params_dtype,
                                     quant_config=None)
        self.experts = nn.ModuleList([
            FeedForward(hidden_size=self.hidden_size,
                        intermediate_size=self.intermediate_size,
                        hidden_act=self.hidden_act,
                        up_proj_name=self.up_proj_name,
                        is_gated=self.is_gated,
                        down_proj_name=self.down_proj_name,
                        bias=self.has_bias,
                        quant_config=self.quant_config,
                        skip_bias_add=self.skip_bias_add,
                        reduce_results=False) for idx in range(self.num_experts_per_rank)
        ])

        self.init_pack_param()


    def init_pack_param(self):
        self.w13 = None
        self.w2 = None
        self.b13 = None
        self.b2 = None
        self.w13_scale = None
        self.w2_scale = None
        self.a13_scale = None
        self.a2_scale = None
        self.pack_params_done = False


    def map_param_data(self, param_list, is_use_first_data=False):
        if len(param_list) == 0:
            return None

        if is_use_first_data or len(param_list) == 1:
            first_data = param_list[0].data
            for param in param_list[1: -1]:
                param.data = first_data
            out_param = first_data.view_as(param_list[0])
        else:
            packed_param = torch._utils._flatten_dense_tensors(param_list)
            data_list = torch._utils._unflatten_dense_tensors(packed_param, param_list)
            for data, param in zip(data_list, param_list):
                param.data = data
            out_param = packed_param.view(len(param_list), *data_list[0].shape)

        torch.mlu.empty_cache()

        return out_param


    def pack_unquantized_params(self, w13, w2, b13, b2):
        for expert in self.experts:
            up_proj = getattr(expert, self.up_proj_name)
            down_proj = getattr(expert, self.down_proj_name)
            w13.append(up_proj.weight)
            w2.append(down_proj.weight)
            if self.has_bias:
                b13.append(up_proj.bias)
                b2.append(down_proj.bias)


    def pack_smoothquant_params(self, w13, w2, b13, b2, w13_scale, w2_scale, a13_scale, a2_scale):
        for expert in self.experts:
            up_proj = getattr(expert, self.up_proj_name)
            down_proj = getattr(expert, self.down_proj_name)
            w13.append(up_proj.qweight)
            w2.append(down_proj.qweight)
            if self.has_bias:
                b13.append(up_proj.bias)
                b2.append(down_proj.bias)
            w13_scale.append(up_proj.per_channel_scale)
            w2_scale.append(down_proj.per_channel_scale)
            if self.quant_config.input_quant_method == "per_token":
                a13_scale.append(up_proj.smooth)
                a2_scale.append(down_proj.smooth)
            else:
                a13_scale.append(up_proj.scale_to_int)
                a2_scale.append(down_proj.scale_to_int)


    def pack_weightonly_params(self, w13, w2, b13, b2, w13_scale, w2_scale):
        for expert in self.experts:
            up_proj = getattr(expert, self.up_proj_name)
            down_proj = getattr(expert, self.down_proj_name)
            w13.append(up_proj.qweight)
            w2.append(down_proj.qweight)
            if self.has_bias:
                b13.append(up_proj.bias)
                b2.append(down_proj.bias)
            w13_scale.append(up_proj.per_channel_scale)
            w2_scale.append(up_proj.per_channel_scale)


    def pack_params(self):
        if self.pack_params_done:
            return

        w13 = []
        w2 = []
        b13 = []
        b2 = []
        w13_scale = []
        w2_scale = []
        a13_scale = []
        a2_scale = []

        if self.quant_config is None:
            self.pack_unquantized_params(w13, w2, b13, b2)
        elif isinstance(self.quant_config, SmoothQuantConfig):
            self.pack_smoothquant_params(w13, w2, b13, b2, w13_scale, w2_scale, a13_scale, a2_scale)
        elif isinstance(self.quant_config, WeightOnlyConfig):
            self.pack_weightonly_params(w13, w2, b13, b2, w13_scale, w2_scale)
        else:
            raise ValueError(f'Unsupported quantization:{self.quant_config}')

        # pack weigth
        self.w13 = self.map_param_data(w13)
        self.w2 = self.map_param_data(w2)

        # pack bias
        if self.has_bias:
            self.b13 = self.map_param_data(b13)
            # NOTE: The bias for fc2 is only applied on tp_rank 0. If we added it on all nodes the allreduce() would
            # contain multiple copies of the bias. The bias on other node will be ignored, and may be set to nullptr
            if self.skip_bias_add is False:
                self.b2 = self.map_param_data(b2)


        # pack weight scale
        if len(w13_scale) > 0:
            self.w13_scale = self.map_param_data(w13_scale)
        if len(w2_scale) > 0:
            self.w2_scale = self.map_param_data(w2_scale)

        # pack activate scale
        if len(a13_scale) > 0:
            self.a13_scale = self.map_param_data(a13_scale)
        if len(a2_scale) > 0:
            self.a2_scale = self.map_param_data(a2_scale)

        self.pack_params_done = True


    def forward(self, hidden_states: torch.Tensor, residual: Optional[torch.Tensor] = None) -> torch.Tensor:
        orig_hidden_states_shape = hidden_states.shape
        hidden_states = hidden_states.view(-1, self.hidden_size)
        # expert_logits: [num_tokens, self.num_experts_per_rank]
        expert_logits, _ = self.gate(hidden_states)
        final_hidden_states = self.forward_experts(hidden_states, expert_logits, residual)

        if self.tp_size > 1:
            final_hidden_states = tensor_model_parallel_all_reduce(
                final_hidden_states)

        output = final_hidden_states.view(orig_hidden_states_shape)
        return output


    def forward_experts(self, hidden_states, expert_logits, residual: Optional[torch.Tensor] = None):
        residual_ = None if self.tp_rank > 0 else residual
        if self.is_use_fused_moe and self.expert_group != 1:
            final_hidden_states = self.forward_group_experts(hidden_states, expert_logits, residual_)
        elif self.is_use_fused_moe:
            self.pack_params()
            final_hidden_states = mlu_ops.fused_moe(hidden_states=hidden_states,
                                                    gating_output=expert_logits,
                                                    w1=self.w13,
                                                    w2=self.w2,
                                                    bias1=self.b13,
                                                    bias2=self.b2,
                                                    residual=residual_,
                                                    input_smooth=self.a13_scale,
                                                    act_smooth=self.a2_scale,
                                                    w1_scale=self.w13_scale,
                                                    w2_scale=self.w2_scale,
                                                    topk=self.top_k,
                                                    renormalize=self.renormalize,
                                                    gated=self.is_gated,
                                                    act_mode=self.hidden_act,
                                                    start_expert_id=self.start_expert_id)
        else:
            final_hidden_states = self.forward_experts_nofused(hidden_states, expert_logits)
            if residual_ is not None:
                final_hidden_states = final_hidden_states + residual_
        return final_hidden_states


    def forward_experts_nofused(self, hidden_states, expert_logits):
        hidden_states_shape = hidden_states.shape
        topk_values, topk_indices = self.topk_softmax(expert_logits)
        expand_gather_idx, scatter_idx, expand_token_count, cusum_token_count = self.generate_gather_idx(
            topk_indices)
        # no expert is routed, then expand_gather_idx, expand_scatter_idx has no item,
        # expand_token_count and expand_cusum_token_count has item but the value is all zero
        # so this rank should only return final_hidden_states with zero value
        if expand_gather_idx.numel() == 0:
            final_hidden_states = torch.zeros_like(hidden_states,
                                                   dtype=hidden_states.dtype,
                                                   device=hidden_states.device)
            return final_hidden_states

        expand_hidden_states = self.expand_input(hidden_states, expand_gather_idx)

        expand_output_list = []
        expand_cusum_token_count = cusum_token_count[self.start_expert_id:self.end_expert_id +
                                                     1] - cusum_token_count[self.start_expert_id]
        for expert_idx, num_tokens_per_expert in enumerate(expand_token_count):
            if num_tokens_per_expert > 0:
                expert_hidden_states = expand_hidden_states[
                    expand_cusum_token_count[expert_idx]:expand_cusum_token_count[expert_idx + 1]]
                expert_output = self.experts[expert_idx](expert_hidden_states)
                expert_output = expert_output[0] if isinstance(expert_output, (tuple, list)) else expert_output
                expand_output_list.append(expert_output)
        expand_output = torch.cat(expand_output_list, dim=0)
        final_hidden_states = self.combine_moe(expand_output, scatter_idx, cusum_token_count, hidden_states_shape,
                                               topk_values)

        return final_hidden_states



    def forward_group_experts(self, hidden_states, expert_logits, residual_):
        ori_input_shape = hidden_states.shape
        dtype = hidden_states.dtype
        self.pack_params()
        gating_output=expert_logits.to(torch.float32)
        w1=self.w13
        w2=self.w2
        bias1=self.b13
        bias2=self.b2
        input_smooth=self.a13_scale
        act_smooth=self.a2_scale
        w1_scale=self.w13_scale
        w2_scale=self.w2_scale
        topk=self.top_k
        renormalized=self.renormalize
        gated=self.is_gated
        act_mode=self.hidden_act

        start_expert_id=self.start_expert_id
        expert_num = gating_output.size(-1)
        expert_size = w1.size(0)
        max_m = hidden_states.shape[0]
        hidden_states = hidden_states.view(-1, hidden_states.size(-1))
        gating_output = gating_output.view(-1, gating_output.size(-1))
        residual_ = residual_.view(-1, residual_.size(-1)) if residual_ is not None else None
        per_token_sq = False
        # check quant
        check_list = [input_smooth, act_smooth, w1_scale, w2_scale]
        if all(x is not None for x in check_list):
            per_token_sq = True

        if not (all(x is None for x in check_list) or all(x is not None for x in check_list)):
            raise ValueError("input_smooth, act_smooth, w1_scale and w2_scale must be present "
                             "and absent at the same time.")
        # softmax_topk
        reduce_weight, expert_id = mlu_ops.moe_softmax_topk(gating_output,
                                                            topk, renormalized)
        # gen_idx
        expand_idx, combine_idx, token_count, cusum_token_count = mlu_ops.moe_gen_idx(expert_id, expert_num)
        # check quant
        if per_token_sq:
            major, minor = current_platform.get_device_capability()
            if major == 3:
                expand_hidden_states = mlu_ops.moe_expand_input(hidden_states, expand_idx,
                    cusum_token_count, start_expert_id, expert_size)
                quant_input, input_scale = mlu_ops.moe_quantize(expand_hidden_states,
                    input_smooth, None, token_count[start_expert_id:start_expert_id+expert_size])
            else:
                quant_input, input_scale = mlu_ops.moe_quantize(hidden_states,
                    input_smooth, None, token_count[start_expert_id:start_expert_id+expert_size], expand_idx,
                    cusum_token_count[start_expert_id].unsqueeze(0))
        else:
            expand_hidden_states = mlu_ops.moe_expand_input(hidden_states, expand_idx,
                    cusum_token_count, start_expert_id, expert_size)

        if per_token_sq:
            gemm1_out = mlu_ops.smooth_quant_group_gemm(quant_input, w1,
                                                    token_count[start_expert_id:start_expert_id+expert_size],
                                                    None, None, None, None,
                                                    input_scale, w1_scale, dtype, max_m)
        else:
            gemm1_out = mlu_ops.group_gemm(expand_hidden_states, w1,
                                       token_count[start_expert_id:start_expert_id+expert_size],
                                       None, None, None, None, max_m)
        # add_bias_active
        act_out = mlu_ops.moe_active(gemm1_out, act_mode, gated, None, bias1, cusum_token_count, start_expert_id, expert_size)
        if per_token_sq:
            quant_input, input_scale = mlu_ops.moe_quantize(act_out, act_smooth, None,
                                                        token_count[start_expert_id:start_expert_id+expert_size])
        if per_token_sq:
            gemm2_out = mlu_ops.smooth_quant_group_gemm(quant_input, w2,
                                                    token_count[start_expert_id:start_expert_id+expert_size],
                                                    None, None, None, None, input_scale, w2_scale, dtype, max_m)
        else:
            gemm2_out = mlu_ops.group_gemm(act_out, w2,
                                       token_count[start_expert_id:start_expert_id+expert_size],
                                       None, None, None, None, max_m)

        output = mlu_ops.moe_combine_result(gemm2_out, reduce_weight, combine_idx,
                                        residual_, cusum_token_count, start_expert_id,
                                        expert_size, bias2)
        return output.view(ori_input_shape)


    def topk_softmax(self, expert_logits):
        # expert_logits: [num_tokens, self.num_experts_per_rank]
        # topk_values: [num_tokens, self.top_k]
        # topk_indices: [num_tokens, self.top_k]
        if self.renormalize:
            topk_values, topk_indices = torch.topk(expert_logits, self.top_k, dim=-1)
            topk_values = torch.softmax(topk_values, -1)
        else:
            router_probs = torch.softmax(expert_logits, -1)
            topk_values, topk_indices = torch.topk(router_probs, self.top_k, dim=-1)

        return topk_values, topk_indices


    def generate_gather_idx(self, topk_indices):
        device = topk_indices.device
        # gather_expand_idx: [num_tokens * self.top_k]
        sorted_expert_id, indices = topk_indices.flatten().sort()
        gather_idx = indices // self.top_k

        seqs = torch.arange(indices.numel(), dtype=indices.dtype, device=indices.device)
        scatter_idx=torch.zeros((indices.numel(),), dtype=seqs.dtype, device=seqs.device).scatter(0, indices, seqs)

        # token_count: [self.num_experts_per_rank]
        partial_token_index, partial_token_count = sorted_expert_id.unique(sorted=True, return_counts=True)
        zero_token_count = torch.zeros(self.num_total_experts, dtype=partial_token_count.dtype, device=device)
        token_count = zero_token_count.scatter(dim=0, index=partial_token_index, src=partial_token_count)
        # cusum_token_count: [self.num_experts_per_rank + 1]
        cusum_token_count = torch.cat(
            [torch.tensor([0], dtype=token_count.dtype, device=device),
             token_count.cumsum(dim=0)])

        num_tokens_before_expert = cusum_token_count[self.start_expert_id]
        num_tokens_including_expert = cusum_token_count[self.end_expert_id]

        expand_gather_idx = gather_idx[num_tokens_before_expert:num_tokens_including_expert]
        expand_token_count = token_count[self.start_expert_id:self.end_expert_id]

        return expand_gather_idx, scatter_idx, expand_token_count, cusum_token_count


    def expand_input(self, hidden_states, expand_gather_idx):
        expand_hidden_states = hidden_states[expand_gather_idx]
        return expand_hidden_states


    def combine_moe(self, expand_output, scatter_idx, cusum_token_count, hidden_states_shape, topk_values):
        num_tokens, hidden_size = hidden_states_shape
        num_tokens_before_expert = cusum_token_count[self.start_expert_id]
        num_tokens_after_expert = cusum_token_count[-1] - cusum_token_count[self.end_expert_id]

        expand_output_before_expert = torch.zeros((num_tokens_before_expert, hidden_size),
                                                   dtype=expand_output.dtype,
                                                   device=expand_output.device)
        expand_output_after_expert = torch.zeros((num_tokens_after_expert, hidden_size),
                                                   dtype=expand_output.dtype,
                                                   device=expand_output.device)
        unscatted_output = torch.cat([expand_output_before_expert, expand_output, expand_output_after_expert], dim=0)
        scatter_output = unscatted_output[scatter_idx]
        hidden_states_weight = topk_values.flatten().unsqueeze(-1)
        weighted_hidden_states = scatter_output * hidden_states_weight
        unreduced_hidden_states = weighted_hidden_states.view(num_tokens, self.top_k, hidden_size)
        final_hidden_states = unreduced_hidden_states.sum(dim=1)

        return final_hidden_states
