from typing import Callable, Optional
from collections.abc import Iterable

import torch
import torch.distributed as dist

from vllm.logger import init_logger
from vllm.platforms import current_platform
from vllm.config import get_current_vllm_config
from vllm.distributed.parallel_state import get_ep_group
from vllm.model_executor.custom_op import CustomOp
from vllm.forward_context import ForwardContext, get_forward_context
from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig
from vllm.model_executor.layers.quantization.base_config import (
    QuantizationConfig, QuantizeMethodBase)
from vllm.model_executor.layers.fused_moe import FusedMoE
from vllm.model_executor.layers.fused_moe.layer import FusedMoEMethodBase, UnquantizedFusedMoEMethod
from vllm.model_executor.layers.fused_moe.mori_moe.ep_moe_utlis import EpMoeConfig
from vllm.utils import direct_register_custom_op

try:
    import mori
    from lmslim.layers.gemm.int8_utils import (
        per_token_quant_int8)
except ImportError:
    is_mori_available = False

logger = init_logger(__name__)

_MORI_OP = None


@CustomOp.register("unquantized_mori_moe")
class UnquantizedMoriMoeMethod(UnquantizedFusedMoEMethod):
    """MoE method without quantization."""

    def __init__(self, moe: FusedMoEConfig):
        super().__init__(moe)
        self.topk_indices_dtype = None
        self.moe = moe

        self.rocm_aiter_moe_enabled = False

    def apply_mori_ep(
            self,
            layer: torch.nn.Module,
            hidden_states: torch.Tensor,
            topk_weights: torch.Tensor,
            topk_ids: torch.Tensor,
            inplace=True,
            activation: str = "silu",
            apply_router_weight_on_input: bool = False,
            global_num_experts: int = -1,
            expert_map: Optional[torch.Tensor] = None,
            use_nn_moe: Optional[bool] = False,
    ) -> torch.Tensor:

        return self.forward(
            hidden_states=hidden_states,
            layer=layer,
            topk_weights=topk_weights,
            topk_ids=topk_ids,
            inplace=inplace,
            activation=activation,
            apply_router_weight_on_input=apply_router_weight_on_input,
            global_num_experts=global_num_experts,
            expert_map=expert_map,
            use_nn_moe=use_nn_moe)

    def forward_cuda(
            self,
            layer: torch.nn.Module,
            hidden_states: torch.Tensor,
            topk_weights: torch.Tensor,
            topk_ids: torch.Tensor,
            inplace=True,
            activation: str = "silu",
            apply_router_weight_on_input: bool = False,
            global_num_experts: int = -1,
            expert_map: Optional[torch.Tensor] = None,
            use_nn_moe: Optional[bool] = False,
    ) -> torch.Tensor:

        # process MoE
        def custom_forward(layer, hidden_states):
            expert_output = self.fused_experts(
                hidden_states=hidden_states,
                w1=layer.w13_weight,
                w2=layer.w2_weight,
                topk_weights=topk_weights,
                topk_ids=topk_ids,
                inplace=True,
                activation=activation,
                apply_router_weight_on_input=apply_router_weight_on_input,
                global_num_experts=global_num_experts,
                expert_map=expert_map,
                use_nn_moe=use_nn_moe
            )

            return expert_output

        output = custom_forward(layer, hidden_states)

        return output

    def forward_cpu(
            self,
            layer: torch.nn.Module,
            hidden_states: torch.Tensor,
            topk_weights: torch.Tensor,
            topk_ids: torch.Tensor,
            inplace=True,
            activation: str = "silu",
            apply_router_weight_on_input: bool = False,
            global_num_experts: int = -1,
            expert_map: Optional[torch.Tensor] = None,
            use_nn_moe: Optional[bool] = False,
            **kwargs,
    ):
        raise NotImplementedError

    def forward_hpu(
            self,
            layer: torch.nn.Module,
            hidden_states: torch.Tensor,
            topk_weights: torch.Tensor,
            topk_ids: torch.Tensor,
            inplace=True,
            activation: str = "silu",
            apply_router_weight_on_input: bool = False,
            global_num_experts: int = -1,
            expert_map: Optional[torch.Tensor] = None,
            use_nn_moe: Optional[bool] = False,
    ) -> torch.Tensor:
        raise NotImplementedError

    def forward_tpu(
            self,
            layer: torch.nn.Module,
            hidden_states: torch.Tensor,
            topk_weights: torch.Tensor,
            topk_ids: torch.Tensor,
            inplace=True,
            activation: str = "silu",
            apply_router_weight_on_input: bool = False,
            global_num_experts: int = -1,
            expert_map: Optional[torch.Tensor] = None,
            use_nn_moe: Optional[bool] = False,
    ) -> torch.Tensor:
        raise NotImplementedError

    if current_platform.is_tpu():
        forward_native = forward_tpu
    elif current_platform.is_cpu():
        forward_native = forward_cpu
    else:
        forward_native = forward_cuda


class MoriMoE(FusedMoE):
    """
    dp+ep MoE Expert Parallel Impl

    """

    def __init__(
            self,
            num_experts: int,  # Global number of experts
            top_k: int,
            hidden_size: int,
            intermediate_size: int,
            params_dtype: Optional[torch.dtype] = None,
            reduce_results: bool = False,
            renormalize: bool = True,
            use_grouped_topk: bool = False,
            num_expert_group: Optional[int] = None,
            topk_group: Optional[int] = None,
            quant_config: Optional[QuantizationConfig] = None,
            tp_size: Optional[int] = None,
            ep_size: Optional[int] = None,
            dp_size: Optional[int] = None,
            prefix: str = "",
            custom_routing_function: Optional[Callable] = None,
            scoring_func: str = "softmax",
            e_score_correction_bias: Optional[torch.Tensor] = None,
            apply_router_weight_on_input: bool = False,
            activation: str = "silu",
            routed_scaling_factor: Optional[float] = None,
            enable_eplb: bool = False,
            num_redundant_experts: int = 0,
            moe_permute_fusion: bool = False,
    ):
        super().__init__(num_experts, top_k, hidden_size,
                         intermediate_size, params_dtype,
                         reduce_results, renormalize,
                         use_grouped_topk, num_expert_group,
                         topk_group, quant_config, tp_size,
                         ep_size, dp_size, prefix,
                         custom_routing_function, scoring_func,
                         e_score_correction_bias,
                         apply_router_weight_on_input,
                         activation,
                         routed_scaling_factor=routed_scaling_factor,
                         enable_eplb=enable_eplb,
                         num_redundant_experts=num_redundant_experts,
                         )

        self.ep_moe_config: EpMoeConfig = EpMoeConfig.make(
            moe_router_topk=self.top_k,
            # TODO: support fusion permute
            moe_permute_fusion=moe_permute_fusion,
            ep_size=self.ep_size,
            num_moe_experts=self.global_num_experts,
            routed_scaling_factor=self.routed_scaling_factor,
            apply_router_weight_on_input=self.apply_router_weight_on_input
        )

        local_expert_indices_offset = (
                self.ep_rank * self.local_num_experts
        )
        self.local_expert_indices = [
            local_expert_indices_offset + i for i in range(self.local_num_experts)
        ]
        
        self.shared_experts = None

        self.scales = None
        self.use_int8_dispatch = True

        vllm_config = get_current_vllm_config()
        self.max_num_inp_token_per_rank = 1024 #vllm_config.scheduler_config.max_num_seqs
        self.mori_op = self.get_mori_op()

    def get_mori_op(self):
        global _MORI_OP
        if _MORI_OP is None:
            world_group = torch.distributed.group.WORLD
            assert world_group is not None
            torch._C._distributed_c10d._register_process_group("mori_ep", get_ep_group().device_group)
            mori.shmem.shmem_torch_process_group_init("mori_ep")

            vllm_config = get_current_vllm_config()
            multi_node = self.ep_size / 8 > 1

            mori_data_type = vllm_config.model_config.dtype
            mori_scale_type_size = vllm_config.model_config.dtype.itemsize
            if self.use_int8_dispatch:
                mori_scale_type_size = 4

            config = mori.ops.EpDispatchCombineConfig(
                data_type=mori_data_type,
                rank=self.ep_rank,
                world_size=self.ep_size,
                hidden_dim=self.hidden_size,
                scale_dim=1 if self.use_int8_dispatch else 0,
                scale_type_size=mori_scale_type_size,
                max_num_inp_token_per_rank=self.max_num_inp_token_per_rank,
                num_experts_per_rank=self.local_num_experts,
                num_experts_per_token=self.top_k,
                max_token_type_size=2,
                block_num=80,
                warp_num_per_block=4,
                kernel_type=mori.ops.EpDispatchCombineKernelType.InterNode if multi_node else \
                    mori.ops.EpDispatchCombineKernelType.IntraNode
            )
            _MORI_OP = mori.ops.EpDispatchCombineOp(config)

        return _MORI_OP

    def set_shared_experts(self, shared_experts: torch.nn.Module):
        if self.shared_experts is None:
            self.shared_experts = shared_experts

    def create_quant_method(self, moe, quant_config, prefix):
        # Note: get_quant_method will look at the layer's local_num_experts
        # for heuristic purposes, so it must be initialized first.
        quant_method: Optional[QuantizeMethodBase] = None
        quant_method = (UnquantizedMoriMoeMethod(moe) if quant_config is None
                        else quant_config.get_quant_method(self, prefix))

        assert quant_method is not None
        assert isinstance(quant_method, FusedMoEMethodBase)
        return quant_method

    def sync(self):
        # torch.cuda.synchronize()
        dist.barrier()

    def forward(self, hidden_states: torch.Tensor,
                router_logits: torch.Tensor):
        return torch.ops.vllm.mori_moe_forward(hidden_states, router_logits,
                                             self.layer_name)

    def get_expert_weights(self) -> Iterable[torch.Tensor]:
        weights = list(self.named_parameters())

        # Filter out the non-expert weights.
        # `e_score_correction_bias` is a bias for each logical expert,
        # with shape (num_logical_experts,), not an expert weight.
        NON_EXPERT_WEIGHTS = {
            "e_score_correction_bias",
            "shared_experts.gate_up_proj.weight",
            "shared_experts.gate_up_proj.weight_scale",
            "shared_experts.down_proj.weight",
            "shared_experts.down_proj.weight_scale"
        }

        return [
            weight.view(self.local_num_experts, -1) for name, weight in weights
            if name not in NON_EXPERT_WEIGHTS
        ]

    def forward_impl(self, hidden_states: torch.Tensor,
                     router_logits: torch.Tensor):

        topk_weights, topk_ids = self.select_experts(
            hidden_states=hidden_states,
            router_logits=router_logits,
            use_grouped_topk=self.use_grouped_topk,
            top_k=self.top_k,
            renormalize=self.renormalize,
            topk_group=self.topk_group,
            num_expert_group=self.num_expert_group,
            custom_routing_function=self.custom_routing_function,
            scoring_func=self.scoring_func,
            e_score_correction_bias=self.e_score_correction_bias,
            indices_type=torch.int32,
            routed_scaling_factor=self.routed_scaling_factor,
            use_fused_gate=self.use_fused_gate)

        if self.shared_experts is not None:
            shared_output = self.shared_experts(hidden_states)

        if self.use_int8_dispatch:
            hidden_states, scales = per_token_quant_int8(hidden_states)
        else:
            if self.scales is None:
                self.scales = torch.rand(
                    hidden_states.shape[0],
                    0,
                    dtype=torch.float32,
                    device=hidden_states.device,
                )
            scales = self.scales

        # self.sync()

        (
            dispatch_output,
            dispatch_weights,
            dispatch_scales,
            dispatch_indices,
            dispatch_recv_num_token,
        ) = self.mori_op.dispatch(
            hidden_states,
            topk_weights,
            scales,
            topk_ids
        )

        expert_output = self.quant_method.apply_mori_ep(
            layer=self,
            x=dispatch_output,
            topk_weights=dispatch_weights,
            topk_ids=dispatch_indices,
            global_num_experts=self.global_num_experts,
            expert_map=self.expert_map,
            activation=self.activation,
            apply_router_weight_on_input=self.apply_router_weight_on_input,
            use_nn_moe=self.use_nn_moe,
            num_local_tokens=dispatch_recv_num_token,
            expect_m=hidden_states.shape[0],
            scales=dispatch_scales if self.use_int8_dispatch else None
        )
        
        # self.sync()

        combine_output, _ = self.mori_op.combine(expert_output, dispatch_weights, topk_ids)
        final_hidden_states = combine_output[:hidden_states.shape[0], :]

        # self.sync()

        if self.shared_experts is not None:
            if hidden_states.dtype != torch.float16:
                final_hidden_states = final_hidden_states + shared_output
            else:
                # Fix FP16 overflow
                # See DeepseekV2DecoderLayer for more details.
                final_hidden_states = final_hidden_states + shared_output \
                                      * (1. / self.routed_scaling_factor)

        return final_hidden_states


def mori_moe_forward(hidden_states: torch.Tensor, router_logits: torch.Tensor,
                   layer_name: str) -> torch.Tensor:
    forward_context: ForwardContext = get_forward_context()
    self = forward_context.no_compile_layers[layer_name]
    assert self.quant_method is not None

    return self.forward_impl(hidden_states, router_logits)


def mori_moe_forward_fake(hidden_states: torch.Tensor, router_logits: torch.Tensor,
                        layer_name: str) -> torch.Tensor:
    return torch.empty_like(hidden_states)


direct_register_custom_op(
    op_name="mori_moe_forward",
    op_func=mori_moe_forward,
    mutates_args=["hidden_states", "router_logits"],
    fake_impl=mori_moe_forward_fake,
    dispatch_key=current_platform.dispatch_key,
    tags=(torch.Tag.needs_fixed_stride_order,),
)