
import os
from abc import ABC, abstractmethod
from typing import List, Optional, Tuple

import torch
import torch.nn as nn

from vllm.distributed.parallel_state import (get_dp_group,
                                                     get_tp_group,
                                                     get_ep_group,
                                                     get_tensor_model_parallel_rank)
from vllm.model_executor.layers.fused_moe.ep_moe.ep_moe_utlis import (EPSharedExperts, 
                                                                      maybe_move_tensor_to_cpu,
                                                                      maybe_move_tensor_to_cpu_block,
                                                                      permute,
                                                                      sort_chunks_by_idxs,
                                                                      unpermute,
                                                                      all_to_all,
                                                                      EpMoeConfig)
from vllm.distributed import (tensor_model_parallel_all_gather,
                              tensor_model_parallel_gather,
                              expert_parallel_all_gather,
                              expert_parallel_gather)
from vllm.platforms import current_platform
from vllm.forward_context import ForwardContext, get_forward_context
from vllm.utils import direct_register_custom_op
from vllm.config import get_current_vllm_config
from lightop import groupgemm_permute, groupgemm_unpermute


cuda_dtoh_stream = torch.cuda.Stream()
cuda_dtoh_sync_event = torch.cuda.Event(enable_timing=False)


class MoETokenDispatcher(nn.Module):
    """
    MoE Token Dispatcher
    """

    def __init__(self, config: EpMoeConfig) -> None:
        """
        Initialize the MoE Token Dispatcher.
        """
        super().__init__()
        self.config = config

        self.tp_size = 1
        self.ep_size = config.ep_size

    @property
    def ep_group(self):
        """Get expert model parallel group."""
        return get_ep_group()

    @property
    def tp_group(self):
        """Get expert tensor parallel group."""
        return get_tp_group()

    @property
    def tp_rank(self):
        """Get expert tensor parallel rank."""
        return 0#get_tensor_model_parallel_rank()

    @property
    def tp_ep_group(self):
        """Get expert tensor and model parallel group."""
        return get_ep_group()

    @abstractmethod
    def token_permutation(
        self, tokens: torch.Tensor, probs: torch.Tensor, routing_map: torch.Tensor
    ):
        """Dispatch tokens to experts.

        Args:
            tokens (torch.Tensor): Input tokens.
            probs (torch.Tensor): The routing probability tensor [num_tokens, num_experts].
            routing_map (torch.Tensor): Token to expert mapping tensor.

        Returns:
            torch.Tensor: Tokens tensor.
        """
        raise NotImplementedError("Dispatch function not implemented.")

    @abstractmethod
    def token_unpermutation(self, expert_output: torch.Tensor, bias: torch.Tensor = None):
        """Restores the expert output to its original ordering.

        Args:
            expert_output (torch.Tensor): The output tensor from the expert models.
            bias (torch.Tensor): The bias tensor.

        Returns:
            (torch.Tensor, torch.Tensor): Unpermuted activation and optional bias.
        """
        raise NotImplementedError("Restore function not implemented.")

    def set_shared_experts(self, shared_experts):
        """Set shared expert to the dispatcher."""
        assert self.config.moe_shared_expert_overlap
        self.shared_experts = shared_experts


class MoEAlltoAllTokenDispatcher(MoETokenDispatcher):
    """
    AlltoAll-based token dispatcher.

    The workflow of AlltoAll token dispatcher is as follows:
    (1) preprocess(): calculate necessary metadata for communication and permute
    (2) token_permutation(): permute->A2A(EP)->AG(TP)->sort_chunk(if num_local_experts>1)
    (3) token_unpermutation(): sort_chunk(if num_local_experts>1)->RS(TP)->A2A(EP)->unpermute
    """

    def __init__(
        self, num_local_experts: int, local_expert_indices: List[int], config: EpMoeConfig, layer_name: str=""
    ) -> None:
        """
        Initialize the AlltoAll token dispatcher.

        Args:
            num_local_experts (int): Number of local experts on the current device.
            local_expert_indices (List[int]): Indices of local experts on the current device.
            config (TransformerConfig): Configuration for the transformer model.
        """
        super().__init__(config=config)
        self.num_local_experts = num_local_experts
        assert config.num_moe_experts is not None
        self.num_experts = config.num_moe_experts
        assert self.num_local_experts > 0, "Expected at least one expert"
        self.local_expert_indices = local_expert_indices
        assert (
            len(self.local_expert_indices) == self.num_local_experts
        ), "Invalid local expert indices"
        for i in range(len(self.local_expert_indices) - 1):
            assert (
                self.local_expert_indices[i] == self.local_expert_indices[i + 1] - 1
            ), "local_expert_indices must be continous"

        self.layer_name = layer_name
        # [ep_size]. Represents the number of tokens sent by the current rank to other
        # EP ranks.
        self.input_splits = None
        # [ep_size]. Represents the number of tokens received by the current rank from
        # other EP ranks.
        self.output_splits = None
        # [tp_size]. Represents the number of tokens received by the current rank from
        # other TP ranks.
        #self.output_splits_tp = None
        self.permute_idx_device = torch.device("cuda") if self.config.moe_permute_fusion else None
        input_chunk_idxs = torch.arange(
            self.num_experts * self.tp_size, device=self.permute_idx_device
        )
        # [num_local_experts, tp_size * ep_size]. Sort the input chunks by local experts.
        self.sort_input_by_local_experts = input_chunk_idxs.reshape(
            -1, self.num_local_experts
        ).T.ravel()
        # [tp_size * ep_size, num_local_experts]. Restore the output chunks by local experts.
        self.restore_output_by_local_experts = input_chunk_idxs.reshape(
            self.num_local_experts, -1
        ).T.ravel()

        # A cuda stream synchronization is needed in self.token_permutation() in some cases,
        # because there are several non-blocking DtoH data transfers called at
        # `self.cuda_dtoh_point`. The synchronization happens at `self.cuda_sync_point`, which is
        # decided based on the MoE and parallel settings. Valid points are "before_permutation_1",
        # "before_ep_alltoall", "before_permutation_2", "before_finish", and "no_sync".
        self.cuda_sync_point = "no_sync"
        self.cuda_sync_point_priority = {
            "before_permutation_1": 0,
            "before_ep_alltoall": 1,
            "before_permutation_2": 2,
            "before_finish": 3,
            "no_sync": 4,
        }
        self.cuda_dtoh_point = "before_permutation_1"
        #self.cuda_dtoh_stream = torch.cuda.Stream()
        
        # Whether to use gather or all-gather to gather the logits.
        self.use_all_gather = current_platform.use_all_gather()

        self.probs = None

        # For smuggling this layer into the fused moe custom op
        vllm_config = get_current_vllm_config()
        compilation_config = vllm_config.compilation_config
        if layer_name in compilation_config.static_forward_context:
            raise ValueError("Duplicate layer name: {}".format(layer_name))
        compilation_config.static_forward_context[layer_name] = self

    def preprocess(self, routing_map: torch.Tensor) -> torch.Tensor:
        """
        Preprocess token routing map for AlltoAll communication and token permutation.

        This method computes the number of tokens assigned to each expert based on the routing_map.
        It also initializes the necessary data structures for AlltoAll communication, such as input
        and output splits, and the mapping between global tokens and local experts. This method
        should not call any DtoH data copying due to performance consideration. The necessary DtoH
        copies are made on the `self.cuda_dtoh_stream` at `self.cuda_dtoh_point`.

        Args:
            routing_map (torch.Tensor): The mapping of tokens to experts, with shape
                [num_tokens, num_experts].

        Returns:
            torch.Tensor: Tensor containing the number of tokens assigned to local expert.
        """
        # [num_experts], number of tokens assigned to each expert from the current rank's input.
        num_local_tokens_per_expert = routing_map.sum(dim=0).long()

        self.num_out_tokens = routing_map.size(0) * self.config.moe_router_topk

        # ===================================================
        # Calculate input_splits, output_splits for alltoall/allgather in variable size.
        # ===================================================
        # [ep_size]. Represents the number of tokens sent by the current rank to other
        # EP ranks.
        self.input_splits = num_local_tokens_per_expert.reshape(
            self.ep_size, self.num_local_experts
        ).sum(axis=1)
        # Gather the global distribution of tokens across ranks.
        # num_global_tokens_per_expert represents the number of tokens sent to each
        # expert by all ranks.
        # [tp_size, ep_size, num_experts]
        if self.use_all_gather:
            # Gather is not supported for some devices such as TPUs.
            # Use all-gather instead.
            num_global_tokens_per_expert = expert_parallel_all_gather(num_local_tokens_per_expert) \
                .reshape(self.ep_size, self.tp_size, self.num_experts) \
                .transpose(0, 1)
        else:
            # None may be returned for rank > 0
            num_global_tokens_per_expert = expert_parallel_gather(num_local_tokens_per_expert) \
                .reshape(self.ep_size, self.tp_size, self.num_experts) \
                .transpose(0, 1)

        # [tp_size, ep_size, num_experts] -> [tp_size, ep_size, num_local_experts]
        num_global_tokens_per_local_expert = num_global_tokens_per_expert[
            :, :, self.local_expert_indices[0] : self.local_expert_indices[-1] + 1
        ].contiguous()
        # [tp_size, ep_size, num_local_experts] -> [tp_size, ep_size]
        num_global_tokens_per_rank = num_global_tokens_per_local_expert.sum(axis=2)
        # [tp_size, ep_size] -> [ep_size]
        # self.output_splits represents the number of tokens received by the current rank
        # from other EP rank.
        self.output_splits = num_global_tokens_per_rank[self.tp_rank]
        # [tp_size, ep_size] -> [tp_size]
        # self.output_splits_tp represents the number of tokens received by the current
        # rank from other TP rank.
        #self.output_splits_tp = num_global_tokens_per_rank.sum(axis=1)
        # [tp_size, ep_size, num_local_experts] -> [num_local_experts]
        num_tokens_per_local_expert = num_global_tokens_per_local_expert.sum(dim=(0, 1))

        # A synchronization is needed before expert parallel AlltoAll communication
        # to get the `input_splits` and `output_splits` CPU values.
        #self._maybe_update_cuda_sync_point("before_ep_alltoall")

        if self.num_local_experts > 1:
            # [tp_size * ep_size, num_local_experts]. Represents the number of tokens sent
            # to each local expert by all ranks.
            self.num_global_tokens_per_local_expert = num_global_tokens_per_local_expert.view(
                -1, self.num_local_experts
            )
            # if not self.config.moe_permute_fusion:
            #     # A synchronization is needed before permutation 2
            #     # to get the `num_global_tokens_per_local_expert` CPU value.
            #     self._maybe_update_cuda_sync_point("before_permutation_2")

        # assert (
        #     self.cuda_sync_point_priority[self.cuda_dtoh_point]
        #     <= self.cuda_sync_point_priority[self.cuda_sync_point]
        # ), "cuda_sync_point must be after cuda_dtoh_point."
        return num_tokens_per_local_expert

    def token_permutation(
        self, hidden_states: torch.Tensor,
        probs: torch.Tensor,
        routing_map: torch.Tensor,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        self.routing_map = routing_map
        assert routing_map.dim() == 2, "Expected 2D tensor for token2expert mask"
        assert routing_map.dtype == torch.bool, "Expected bool tensor for mask"
        tokens_per_expert = self.preprocess(self.routing_map)

        if self.config.moe_shared_expert_overlap and self.shared_experts is not None:
            self.shared_experts.pre_forward_comm(hidden_states.view(self.hidden_shape))
        
        global_input_tokens = torch.ops.vllm.token_permutation_forward(tokens_per_expert, hidden_states, 
                                                                       probs, routing_map, self.layer_name)
        return global_input_tokens, tokens_per_expert
    
    def token_permutation_impl(
        self, 
        tokens_per_expert: torch.Tensor,
        hidden_states: torch.Tensor,
        probs: torch.Tensor,
        routing_map: torch.Tensor,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Dispatch tokens to local experts using AlltoAll communication.

        This method performs the following steps:
        1. Preprocess the routing map to get metadata for communication and permutation.
        2. Permute input tokens for AlltoAll communication.
        3. Perform expert parallel AlltoAll communication.
        4. Sort tokens by local expert (if multiple local experts exist).

        Args:
            hidden_states (torch.Tensor): Input token embeddings.
            probs (torch.Tensor): The probabilities of token to experts assignment.
            routing_map (torch.Tensor): The mapping of token to experts assignment.

        Returns:
            Tuple[torch.Tensor, torch.Tensor]:
                - Permuted token embeddings for local experts.
                - Number of tokens per expert.
        """
        # Preprocess: Get the metadata for communication, permutation and computation operations.
        # Permutation 1: input to AlltoAll input
        tokens_per_expert = self._maybe_dtoh_and_synchronize(
            "before_permutation_1", tokens_per_expert
        )
        self.hidden_shape = hidden_states.shape
        if self.config.apply_router_weight_on_input:
            self.probs = probs
            assert probs.dim() == 2, "Expected 2D tensor for probs"
        
        hidden_states = hidden_states.view(-1, self.hidden_shape[-1])

        self.hidden_shape_before_permute = hidden_states.shape

        if False:
            permutated_local_input_tokens, self.reversed_local_input_permutation_mapping = permute(
                hidden_states,
                routing_map,
                num_out_tokens=self.num_out_tokens,
                fused=self.config.moe_permute_fusion
            )
        else:
            cuda_permute_result = groupgemm_permute(hidden_states, routing_map)
            
            permutated_local_input_tokens, self.reversed_local_input_permutation_mapping, \
                self.expert_m_count = cuda_permute_result

        # Perform expert parallel AlltoAll communication
        # tokens_per_expert = self._maybe_dtoh_and_synchronize(
        #     "before_ep_alltoall", tokens_per_expert
        # )

        ###test##############
        #cuda_dtoh_stream.synchronize()
        cuda_dtoh_sync_event.synchronize()
        ###test##############

        global_input_tokens = all_to_all(
            self.ep_group.device_group, permutated_local_input_tokens, self.output_splits, self.input_splits
        )

        if self.config.moe_shared_expert_overlap and self.shared_experts is not None:
            self.shared_experts.linear_fc1_forward_and_act(global_input_tokens)

        # Permutation 2: Sort tokens by local expert.
        # tokens_per_expert = self._maybe_dtoh_and_synchronize(
        #     "before_permutation_2", tokens_per_expert
        # )
        if self.num_local_experts > 1:
            global_input_tokens = sort_chunks_by_idxs(
                    global_input_tokens,
                    self.num_global_tokens_per_local_expert.ravel(),
                    self.sort_input_by_local_experts,
                    fused=self.config.moe_permute_fusion,
                )

        #tokens_per_expert = self._maybe_dtoh_and_synchronize("before_finish", tokens_per_expert)

        return global_input_tokens

    def token_unpermutation(
        self, hidden_states: torch.Tensor,
    ) -> torch.Tensor:
        return torch.ops.vllm.token_unpermutation_forward(hidden_states, self.layer_name)
    
    def token_unpermutation_impl(
        self, hidden_states: torch.Tensor,
    ) -> torch.Tensor:
        """
        Reverse the token permutation to restore the original order.

        This method performs the following steps:
        1. Unsort tokens by local expert (if multiple local experts exist).
        2. Perform expert parallel AlltoAll communication to restore the original order.
        3. Unpermute tokens to restore the original order.

        Args:
            hidden_states (torch.Tensor): Output from local experts.
            bias (torch.Tensor, optional): Bias tensor (not supported).

        Returns:
            Tuple[torch.Tensor, Optional[torch.Tensor]]:
                - Unpermuted token embeddings in the original order.
                - None (bias is not supported).
        """
        # Unpermutation 2: Unsort tokens by local expert.
        if self.num_local_experts > 1:
            hidden_states = sort_chunks_by_idxs(
                hidden_states,
                self.num_global_tokens_per_local_expert.T.ravel(),
                self.restore_output_by_local_experts,
                fused=self.config.moe_permute_fusion,
            )

        # Perform expert parallel AlltoAll communication
        # hidden_states: [SEQL, H] -> [SEQL, H/TP]
        permutated_local_input_tokens = all_to_all(
            self.ep_group.device_group, hidden_states, self.input_splits, self.output_splits
        )

        if self.config.moe_shared_expert_overlap and self.shared_experts is not None:
            self.shared_experts.linear_fc2_forward(permutated_local_input_tokens)
            self.shared_experts.post_forward_comm()

        # Unpermutation 1: AlltoAll output to output
        if False:
            output = unpermute(
                permutated_local_input_tokens,
                self.reversed_local_input_permutation_mapping,
                restore_shape=self.hidden_shape_before_permute,
                probs=self.probs,
                routing_map=self.routing_map,
                fused=self.config.moe_permute_fusion,
            )
        else:
            output = groupgemm_unpermute(permutated_local_input_tokens,
                                self.reversed_local_input_permutation_mapping,
                                list(self.hidden_shape_before_permute),
                                self.probs,
                                self.routing_map,
                                self.expert_m_count)

        # Reshape the output tensor
        output = output.view(self.hidden_shape)

        # Add shared experts output
        if self.config.moe_shared_expert_overlap and self.shared_experts is not None:
            shared_output = self.shared_experts.get_output()

            if hidden_states.dtype != torch.float16:
                output = output + shared_output
            else:
                # Fix FP16 overflow
                # See DeepseekV2DecoderLayer for more details.
                output = output + shared_output \
                    * (1. / self.config.routed_scaling_factor)
        return output

    def _maybe_update_cuda_sync_point(self, point: str):
        """
        Update the CUDA sync point if the priority of the new point is higher than the current
        sync point, which means the new point is reached earlier than the current sync point.
        """
        if (
            self.cuda_sync_point_priority[point]
            < self.cuda_sync_point_priority[self.cuda_sync_point]
        ):
            self.cuda_sync_point = point

    def _maybe_dtoh_and_synchronize(
        self, point: str, tokens_per_expert: torch.Tensor = None
    ) -> torch.Tensor:
        """
        Move all possible GPU tensors to CPU and make a synchronization at the expected point.
        """
        if point == self.cuda_dtoh_point:
            # Move all possible GPU tensors to CPU at self.cuda_dtoh_point.
            on_side_stream = torch.cuda.current_stream() != cuda_dtoh_stream
            if on_side_stream:
                cuda_dtoh_stream.wait_stream(torch.cuda.current_stream())
            with torch.cuda.stream(cuda_dtoh_stream):
                # TODO: use MemcpyBatchAsync instead.
                # tokens_per_expert = maybe_move_tensor_to_cpu(
                #     tokens_per_expert, record_stream=on_side_stream
                # )
                self.input_splits = maybe_move_tensor_to_cpu(
                    self.input_splits, as_numpy=True, record_stream=on_side_stream
                )
                self.output_splits = maybe_move_tensor_to_cpu(
                    self.output_splits, as_numpy=True, record_stream=on_side_stream
                )
                # self.output_splits_tp = maybe_move_tensor_to_cpu(
                #     self.output_splits_tp, as_numpy=True, record_stream=on_side_stream
                # )
                self.num_out_tokens = maybe_move_tensor_to_cpu(
                    self.num_out_tokens, record_stream=on_side_stream
                )
                if self.num_local_experts > 1 and not self.config.moe_permute_fusion:
                    self.num_global_tokens_per_local_expert = maybe_move_tensor_to_cpu(
                        self.num_global_tokens_per_local_expert, record_stream=on_side_stream
                    )

                cuda_dtoh_sync_event.record()

            # if point == self.cuda_sync_point:
            #     # Synchronize with the dtoh stream at self.cuda_sync_point.
            #     cuda_dtoh_stream.synchronize()

        return tokens_per_expert
    

def token_permutation_forward(tokens_per_expert: torch.Tensor,
                              hidden_states: torch.Tensor,
                              probs: torch.Tensor,
                              routing_map: torch.Tensor,
                              layer_name: str) -> torch.Tensor:
    forward_context: ForwardContext = get_forward_context()
    self = forward_context.no_compile_layers[layer_name]

    return self.token_permutation_impl(tokens_per_expert, hidden_states, probs, routing_map)


def token_permutation_forward_fake(tokens_per_expert: torch.Tensor,
                                    hidden_states: torch.Tensor,
                                    probs: torch.Tensor,
                                    routing_map: torch.Tensor,
                                    layer_name: str) -> torch.Tensor:
    return torch.empty_like(hidden_states)


direct_register_custom_op(
    op_name="token_permutation_forward",
    op_func=token_permutation_forward,
    mutates_args=["tokens_per_expert", "hidden_states", "probs", "routing_map"],
    fake_impl=token_permutation_forward_fake,
    dispatch_key=current_platform.dispatch_key,
    tags=(torch.Tag.needs_fixed_stride_order, ),
)


def token_unpermutation_forward(hidden_states: torch.Tensor,
                              layer_name: str) -> torch.Tensor:
    forward_context: ForwardContext = get_forward_context()
    self = forward_context.no_compile_layers[layer_name]

    return self.token_unpermutation_impl(hidden_states)


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


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