import os
import logging
from typing import Callable, List, Optional, Tuple
from dataclasses import dataclass

import torch
import torch.nn.functional as F

from vllm.logger import init_logger
from vllm.platforms import current_platform
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.ep_moe.token_dispatcher import MoEAlltoAllTokenDispatcher
from vllm.model_executor.layers.fused_moe.ep_moe.ep_moe_utlis import EpMoeConfig
from vllm.utils import direct_register_custom_op


logger = init_logger(__name__)


@CustomOp.register("unquantized_ep_moe")
class UnquantizedEPGroupedGemmMethod(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 # is_rocm_aiter_moe_enabled()

    def apply(
        self,
        layer: torch.nn.Module,
        hidden_states: torch.Tensor,
        tokens_per_expert: torch.Tensor,
    ) -> torch.Tensor:

        return self.forward(
            hidden_states=hidden_states,
            layer=layer,
            tokens_per_expert=tokens_per_expert)

    def forward_cuda(
        self,
        layer: torch.nn.Module,
        hidden_states: torch.Tensor,
        tokens_per_expert: torch.Tensor,
    ) -> torch.Tensor:

        # process MoE
        def custom_forward(layer, hidden_states, tokens_per_expert):
            tokens_per_expert = tokens_per_expert.cpu().numpy()

            outputs = []
            start_idx = 0
            for i, num_tokens in enumerate(tokens_per_expert):
                end_idx = start_idx + num_tokens
                if num_tokens == 0:
                    continue
                w1 = layer.w13_weight[i]
                w2 = layer.w2_weight[i]
                tokens_for_this_expert = hidden_states[start_idx:end_idx]

                gateup_output = torch.matmul(tokens_for_this_expert, w1.T)

                # Act
                down_input = torch.zeros(
                    gateup_output.shape[0],
                    gateup_output.shape[1] // 2,
                    device=gateup_output.device,
                    dtype=hidden_states.dtype
                )
                torch.ops._C.silu_and_mul(down_input,
                                gateup_output.view(-1, w1.shape[0]))

                expert_out = torch.matmul(down_input, w2.T)

                outputs.append(expert_out)
                start_idx = end_idx

            if len(outputs) > 0:
                expert_output = torch.cat(outputs, dim=0)
            else:
                assert hidden_states.numel() == 0, f"sorted_tokens: should be empty, but got {hidden_states.shape}"
                expert_output = hidden_states

            return expert_output

        output = custom_forward(layer, hidden_states, tokens_per_expert)

        return output

    def forward_cpu(
        self,
        layer: torch.nn.Module,
        hidden_states: torch.Tensor,
        tokens_per_expert: torch.Tensor,
        **kwargs,
    ):
        raise NotImplementedError

    def forward_hpu(
        self,
        layer: torch.nn.Module,
        hidden_states: torch.Tensor,
        tokens_per_expert: torch.Tensor,
    ) -> torch.Tensor:
        raise NotImplementedError

    def forward_tpu(
        self,
        layer: torch.nn.Module,
        hidden_states: torch.Tensor,
        tokens_per_expert: torch.Tensor,
    ) -> 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 EPMoE(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,
        moe_permute_fusion: bool = True,
        moe_shared_expert_overlap: 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
                         )
        
        self.ep_moe_config: EpMoeConfig = EpMoeConfig.make(
            moe_router_topk=self.top_k,
            # TODO: support fusion permute
            moe_permute_fusion=moe_permute_fusion,
            moe_shared_expert_overlap=moe_shared_expert_overlap,
            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.use_shared_expert = False
        self.token_dispatcher = MoEAlltoAllTokenDispatcher(
                self.local_num_experts, self.local_expert_indices, config=self.ep_moe_config
            )
        
        self.shared_expert_overlap = moe_shared_expert_overlap
        self.shared_experts = None

        self.dpsk_fp16_quick = os.environ.get('DPSK_FP16_QUICK') == '1'

    def set_shared_experts(self, shared_experts: torch.nn.Module):
        if self.shared_experts is None:
            self.shared_experts = shared_experts
            if self.shared_expert_overlap:
                self.token_dispatcher.set_shared_experts(self.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 = (UnquantizedEPGroupedGemmMethod(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 forward(self, hidden_states: torch.Tensor,
                router_logits: torch.Tensor):
        return torch.ops.vllm.ep_moe_forward(hidden_states, router_logits,
                                              self.layer_name)        
    
    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.int64,
                routed_scaling_factor=self.routed_scaling_factor,
                use_fused_gate=self.use_fused_gate)
        if not self.ep_moe_config.moe_shared_expert_overlap and self.shared_experts is not None:
            shared_output = self.shared_experts(hidden_states)

        probs = torch.zeros_like(router_logits, dtype=topk_weights.dtype).scatter(1, topk_ids, topk_weights)
        routing_map = torch.zeros_like(router_logits).int().scatter(1, topk_ids, 1).bool()

        (dispatched_input, tokens_per_expert) = self.token_dispatcher.token_permutation(
            hidden_states, probs, routing_map
        )

        # Matrix multiply.
        expert_output = self.quant_method.apply(
            layer=self,
            hidden_states=dispatched_input,
            tokens_per_expert=tokens_per_expert
        )

        final_hidden_states = self.token_dispatcher.token_unpermutation(expert_output)
        if not self.ep_moe_config.moe_shared_expert_overlap and self.shared_experts is not None:
            # if shared_expert_overlap is True, the expert calculation happens in
            # the token_dispatcher to overlap communications and computations
            shared_output = (
                    self.maybe_all_reduce_tensor_model_parallel(
                        shared_output))

            if hidden_states.dtype != torch.float16 or self.dpsk_fp16_quick:
                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 ep_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 ep_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="ep_moe_forward",
    op_func=ep_moe_forward,
    mutates_args=["hidden_states"],
    fake_impl=ep_moe_forward_fake,
    dispatch_key=current_platform.dispatch_key,
    tags=(torch.Tag.needs_fixed_stride_order, ),
)