import math
import torch
import torch.distributed as dist
from typing import Optional, Tuple, Union

def get_global_seed() -> int:
    return 0


def dispatch(x: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]],
             topk_idx: torch.Tensor, topk_weights: Optional[torch.Tensor],
             num_max_tokens_per_rank: int, num_experts: int):
    """
    The reference implementation of dispatching tokens to experts across multiple ranks.

    Not expanded. Sorted by rank and then by token within each rank (i.e. sorted by `src_token_global_idx`).

    Arguments:
        x: input tokens, `[num_tokens, hidden]` or (`[num_tokens, hidden]`, `[num_tokens, num_sf_packs]`).
        topk_idx: top-k expert indices for each token, `[num_tokens, num_topk]`.
        topk_weights: top-k weights for each token, can be None, `[num_tokens, num_topk]`.
        num_max_tokens_per_rank: the maximum number of tokens per rank, must be `>=` the actual number of
            tokens per rank and aligned with DeepEP's `num_max_tokens_per_rank`, since we calculate
            `src_token_global_idx = src_rank_idx * num_max_tokens_per_rank + src_token_local_idx`.
        num_experts: the total number of experts across all ranks.

    Returns:
        recv_x, recv_topk_idx, recv_topk_weights: received tokens, top-k indices and top-k weights after
            dispatching. Out of range `recv_topk_idx` (i.e. that expert is not on the current rank) are
            set to -1.
        recv_src_token_idx: received `src_token_global_idx` for each received token.
        num_recv_tokens_per_rank: number of received tokens from each rank, `[num_ranks]`.
    """
    rank_idx = dist.get_rank()
    num_ranks = dist.get_world_size()

    assert num_experts % num_ranks == 0
    num_experts_per_rank = num_experts // num_ranks

    # Unpack SF
    use_fp8 = isinstance(x, tuple)
    x, sf = x if use_fp8 else (x, None)

    num_tokens, hidden = x.size()
    num_tokens_, num_topk = topk_idx.size()
    assert num_tokens == num_tokens_
    num_sf_packs = 0
    if sf is not None:
        num_tokens__, num_sf_packs = sf.size()
        assert num_tokens == num_tokens__
    if topk_weights is not None:
        num_tokens__, num_topk_ = topk_weights.size()
        assert num_tokens == num_tokens__
        assert num_topk == num_topk_

    # Prepare per-peer send buffers
    send_x_list = []
    send_sf_list = []
    send_topk_idx_list = []
    send_topk_weights_list = []
    send_src_token_idx_list = []
    num_send_tokens_per_rank = torch.zeros((num_ranks, ), dtype=torch.int, device=x.device)
    for dst_rank_idx in range(num_ranks):
        expert_start_idx = dst_rank_idx * num_experts_per_rank
        expert_end_idx = expert_start_idx + num_experts_per_rank

        # Get the indices of tokens
        mask_to_send = ((expert_start_idx <= topk_idx) & (topk_idx < expert_end_idx)).any(dim=1)
        indices_to_send = mask_to_send.nonzero(as_tuple=True)[0]
        num_send_tokens_per_rank[dst_rank_idx] = indices_to_send.numel()

        # Select the data for tokens
        x_to_send = x[indices_to_send]
        sf_to_send = sf[indices_to_send] if use_fp8 else None
        topk_idx_to_send = topk_idx[indices_to_send]
        topk_weights_to_send = topk_weights[indices_to_send]
        masked_topk_idx = torch.where((expert_start_idx <= topk_idx_to_send) & (topk_idx_to_send < expert_end_idx),
                                      topk_idx_to_send, torch.full_like(topk_idx_to_send, -1))

        send_x_list.append(x_to_send)
        send_sf_list.append(sf_to_send)
        send_topk_idx_list.append(masked_topk_idx)
        send_topk_weights_list.append(topk_weights_to_send)
        send_src_token_idx_list.append(indices_to_send)

    send_x = torch.cat(send_x_list, dim=0)
    send_sf = torch.cat(send_sf_list, dim=0) if use_fp8 else None
    send_topk_idx = torch.cat(send_topk_idx_list, dim=0)
    send_topk_weights = torch.cat(send_topk_weights_list, dim=0)
    send_src_token_idx = torch.cat(send_src_token_idx_list, dim=0).to(torch.int)
    send_src_token_idx += rank_idx * num_max_tokens_per_rank

    # Exchange size
    num_recv_tokens_per_rank = torch.empty((num_ranks, ), dtype=torch.int, device=x.device)
    dist.all_to_all_single(num_recv_tokens_per_rank, num_send_tokens_per_rank)
    num_recv_tokens = int(num_recv_tokens_per_rank.sum().item())

    # Exchange main data
    num_send_tokens_per_rank = num_send_tokens_per_rank.tolist()
    num_recv_tokens_per_rank = num_recv_tokens_per_rank.tolist()
    recv_x = torch.empty((num_recv_tokens, hidden), dtype=x.dtype, device=x.device)
    recv_sf = torch.empty((num_recv_tokens, num_sf_packs), dtype=sf.dtype, device=x.device) if use_fp8 else None
    recv_topk_idx = torch.empty((num_recv_tokens, num_topk), dtype=topk_idx.dtype, device=x.device)
    recv_topk_weights = torch.empty((num_recv_tokens, num_topk), dtype=topk_weights.dtype, device=x.device)
    recv_src_token_idx = torch.empty((num_recv_tokens, ), dtype=torch.int, device=x.device)
    dist.all_to_all_single(recv_x, send_x, num_recv_tokens_per_rank, num_send_tokens_per_rank)
    if use_fp8:
        # gloo has no int16 collective: move the SF bytes (bit-exact, same semantics).
        dist.all_to_all_single(recv_sf.view(torch.uint8), send_sf.contiguous().view(torch.uint8),
                               num_recv_tokens_per_rank, num_send_tokens_per_rank)
    dist.all_to_all_single(recv_topk_idx, send_topk_idx, num_recv_tokens_per_rank, num_send_tokens_per_rank)
    dist.all_to_all_single(recv_topk_weights, send_topk_weights, num_recv_tokens_per_rank, num_send_tokens_per_rank)
    dist.all_to_all_single(recv_src_token_idx, send_src_token_idx, num_recv_tokens_per_rank, num_send_tokens_per_rank)

    # Mask top-k indices
    expert_start_idx = rank_idx * num_experts_per_rank
    expert_end_idx = expert_start_idx + num_experts_per_rank
    mask = (expert_start_idx <= recv_topk_idx) & (recv_topk_idx < expert_end_idx)
    recv_topk_idx = recv_topk_idx - expert_start_idx
    recv_topk_idx.masked_fill_(~mask, -1)

    # Pack SF
    recv_x = (recv_x, recv_sf) if use_fp8 else recv_x

    return (recv_x, recv_topk_idx, recv_topk_weights,
            recv_src_token_idx, torch.tensor(num_recv_tokens_per_rank, dtype=torch.int))


def generate_pre_combine_data(src_token_global_idx: torch.Tensor,
                              num_max_tokens_per_rank: int, num_topk: int, hidden: int) -> torch.Tensor:
    """
    Generate data needed for combine from `src_token_global_idx`.
    Recall that `src_token_global_idx = src_rank_idx * num_max_tokens_per_rank + src_token_local_idx`.
    The generated data (denoted as `y`) of the i-th token has a shape of `[num_topk, hidden]`, with

    `y[j, k] = sin((token_seeds * P % max_seed + 1) / max_seed * (k + 1) + sin(seed))`

    where `P=131071` is a large prime, `token_seeds` is calculated via
    `token_seeds = src_token_global_idx[i] * num_topk + j`, and
    `max_seed = num_ranks * num_max_tokens_per_rank * num_topk`.

    Arguments:
        src_token_global_idx: source token global indices, `[num_tokens]`.

    Returns:
        the generated data, `[num_tokens, num_topk, hidden]`.
    """
    num_ranks = dist.get_world_size()
    token_seeds = (src_token_global_idx.unsqueeze(1) * num_topk +
                   torch.arange(num_topk, device=src_token_global_idx.device).unsqueeze(0))  # [num_tokens, num_topk]
    max_seed = num_ranks * num_max_tokens_per_rank * num_topk
    result = torch.sin(
        (((token_seeds * 131071 % max_seed).float() + 1) / max_seed).unsqueeze(-1) *
        torch.arange(1, hidden + 1, device=src_token_global_idx.device, dtype=torch.float32).broadcast_to(1, 1, hidden) +
        math.sin(float(get_global_seed()))
    )
    return result.to(torch.bfloat16)


def ordered_accumulate(data: torch.Tensor, initial_value: Optional[torch.Tensor] = None) -> torch.Tensor:
    """
    Accumulate `data` in order along the num_topk dimension.

    Arguments:
        data: data to be accumulated, `[num_tokens, num_topk, hidden]`.
        initial_value: initial value for accumulation, `[num_tokens, hidden]`.

    Returns:
        the result, `[num_tokens, hidden]`.
    """
    num_topk = data.shape[1]
    if initial_value is None:
        result = torch.zeros((data.shape[0], data.shape[2]), dtype=torch.float32, device=data.device)
    else:
        result = initial_value.clone()
    for i in range(num_topk):
        result += data[:, i, :].float()
    return result.to(data.dtype)


def combine(y: torch.Tensor, topk_idx: torch.Tensor, num_experts: int,
            bias: Optional[Union[Tuple[torch.Tensor, ...], torch.Tensor]],
            reduce_in_local: bool) -> torch.Tensor:
    """
    The reference implementation of (possibly intra-rank reduced) combining tokens.

    Arguments:
        y: input tokens to be combined, `[num_tokens, num_topk, hidden]`.
        topk_idx: `[num_tokens, num_topk]`.
        num_experts: the total number of experts across all ranks.
        bias: 0, 1 or 2 `[num_tokens, hidden]` final bias added to the output.
        reduce_in_local: whether to do reduction within each rank first.
            - `True`  -> reduce selections landing on the same rank before the global accumulation.
            - `False` -> equivalent to `allow_multiple_reduction` being False.

    Returns:
        the combined result, `[num_tokens, hidden]`.
    """
    num_ranks = dist.get_world_size()
    num_tokens, hidden = y.shape[0], y.shape[2]
    num_topk = y.shape[1]

    def grouped_reduce(data_to_reduce: torch.Tensor, group_id: torch.Tensor) -> torch.Tensor:
        """
        Perform in-place grouped reduction on `data_to_reduce` according to `group_id`.
        The summation within each group is performed in strict order along the `num_topk` dimension.
        The result for each group is stored at the rightmost token of that group, and other tokens are
        set to zero.

        Arguments:
            data_to_reduce: data to be reduced, `[num_tokens, num_topk, hidden]`.
            group_id: group IDs for each token, `[num_tokens, num_topk]`.
        """
        # Shuffle to make tokens with the same group_id contiguous
        group_id, src_indices = torch.sort(group_id, dim=-1, stable=True)
        # transformed_src_indices[i, j] = i * num_topk + src_indices[i, j]
        transformed_src_indices = (
            (src_indices + torch.arange(0, num_tokens, device=y.device).unsqueeze(-1) * num_topk).flatten())
        data_to_reduce = data_to_reduce.view(-1, hidden)[transformed_src_indices].view(num_tokens, num_topk, hidden)
        # Perform segmented reduce within each group
        cur_accum_buf = torch.zeros((num_tokens, hidden), dtype=torch.float32, device=y.device)
        for i in range(num_topk):
            is_segment_break = torch.full((num_tokens, ), True, dtype=torch.bool, device=y.device) \
                if i == num_topk - 1 else group_id[:, i] != group_id[:, i + 1]
            cur_accum_buf += data_to_reduce[:, i, :].float()
            # For one token, if `is_segment_break` is True,
            # save the accumulated value and clear the buffer, otherwise, clear `data_to_reduce[:, i, :]`
            segment_break_token_indices = torch.where(is_segment_break)[0]
            data_to_reduce[segment_break_token_indices, i] = cur_accum_buf[segment_break_token_indices].to(data_to_reduce.dtype)
            cur_accum_buf[segment_break_token_indices] = 0.0
            non_segment_break_token_indices = torch.where(~is_segment_break)[0]
            data_to_reduce[non_segment_break_token_indices, i] = 0.0
        # Unshuffle
        # noinspection PyShadowingNames
        result = torch.empty_like(data_to_reduce)
        result.view(-1, hidden)[transformed_src_indices] = data_to_reduce.view(-1, hidden)
        return result.view(num_tokens, num_topk, hidden)

    num_experts_per_rank = num_experts // num_ranks
    if reduce_in_local:
        y = grouped_reduce(y, topk_idx // num_experts_per_rank)
    bias_sum = None
    if isinstance(bias, tuple):
        assert len(bias) <= 2
        for bias_tensor in bias:
            bias_sum = bias_tensor.float() if bias_sum is None else bias_sum + bias_tensor.float()
    elif bias is not None:
        bias_sum = bias.float()
    result = ordered_accumulate(y, bias_sum)
    return result