# Copyright (c) 2025, Huawei Technologies Co., Ltd.  All rights reserved.
# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.

"""
Fused functions used in the MoE router for NPU

This module provides NPU-compatible implementations of MoE router functions.
Currently, these implementations use small operators (non-fused) from Megatron
as placeholders until NPU-specific fused operators are available.

Precision Notes:
- FP64 is currently not supported.
- Inputs are casted into FP32 when loading from global memory.
- All the math/calculations/accumulations are in FP32.
- "scores" is always in FP32 (match the MCore implementation).
- Only cast to low-precision when necessary and the casting only happens in writing to
  global memory. For example, the gradient is required to have the same dtype as the input.
"""
from typing import Optional, Tuple

import torch


__all__ = [
    "fused_moe_aux_loss",
    "fused_compute_score_for_moe_aux_loss",
    "fused_topk_with_score_function",
]


# ===================== Helper Functions =====================


def _compute_topk_with_group_limited(
    scores: torch.Tensor,
    topk: int,
    num_tokens: int,
    num_experts: int,
    num_groups: Optional[int] = None,
    group_topk: Optional[int] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
    """Perform top-k routing with group-limited selection.

    When using group-limited routing:
    1. Experts are divided into 'num_groups' equal-sized groups
    2. For each token, 'group_topk' groups are selected based on routing scores
       (specifically, the sum of top-2 expert scores within each group)
    3. From these selected groups, 'topk' individual experts are chosen

    Args:
        scores (torch.Tensor): Softmax scores generated by the router.
        topk (int): The number of experts to select for each token.
        num_tokens (int): The number of tokens.
        num_experts (int): The number of experts.
        num_groups (int, optional): Number of groups for routed experts.
        group_topk (int, optional): Number of groups selected for each token.

    Returns:
        Tuple[torch.Tensor, torch.Tensor]: Probs and indices tensor.
    """
    # Organize the experts into groups
    # Select groups based on sum of top-(topk/group_topk) routing scores within each group
    group_scores = (
        scores.view(num_tokens, num_groups, -1).topk(topk // group_topk, dim=-1)[0].sum(dim=-1)
    )
    group_idx = torch.topk(group_scores, k=group_topk, dim=-1, sorted=False)[1]
    group_mask = torch.zeros_like(group_scores)
    group_mask.scatter_(1, group_idx, 1)

    # Mask the experts based on selection groups
    score_mask = (
        group_mask.unsqueeze(-1)
        .expand(num_tokens, num_groups, num_experts // num_groups)
        .reshape(num_tokens, -1)
    )

    masked_scores = scores.masked_fill(~score_mask.bool(), float('-inf'))
    probs, top_indices = torch.topk(masked_scores, k=topk, dim=-1)

    return probs, top_indices


def _compute_topk(
    scores: torch.Tensor,
    topk: int,
    num_tokens: int,
    num_experts: int,
    num_groups: Optional[int] = None,
    group_topk: Optional[int] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
    """Compute the top-k indices for the given scores.

    Args:
        scores (torch.Tensor): The scores tensor.
        topk (int): The number of top-k indices to compute.
        num_tokens (int): The number of tokens.
        num_experts (int): The number of experts.
        num_groups (int, optional): The number of groups to compute the top-k indices for.
        group_topk (int, optional): The number of top-k indices to compute for each group.

    Returns:
        Tuple[torch.Tensor, torch.Tensor]: The top-k indices and the top-k scores.
    """
    if group_topk:
        return _compute_topk_with_group_limited(
            scores=scores,
            topk=topk,
            num_tokens=num_tokens,
            num_experts=num_experts,
            num_groups=num_groups,
            group_topk=group_topk,
        )
    else:
        return torch.topk(scores, k=topk, dim=1)


# ===================== fused_topk_with_score_function =====================


def fused_topk_with_score_function(
    logits: torch.Tensor,
    topk: int,
    use_pre_softmax: bool,
    num_groups: Optional[int],
    group_topk: Optional[int],
    scaling_factor: Optional[float],
    score_function: str,
    expert_bias: Optional[torch.Tensor],
) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    Fused topk with score function router for NPU.
    
    This is a placeholder implementation using small operators (non-fused).
    It provides the same interface as the TransformerEngine version but uses
    PyTorch operations instead of fused CUDA kernels.
    
    This implementation follows Megatron's approach and relies on PyTorch's
    automatic differentiation for gradient computation.
    
    Parameters
    ----------
    logits : torch.Tensor in fp32/bf16/fp16
        Router logits of shape [num_tokens, num_experts]
    topk : int
        Number of experts to select for each token
    use_pre_softmax : bool
        If enabled, the computation order: softmax -> topk.
        Otherwise: topk -> softmax
    num_groups : int, optional
        Used in the group topk
    group_topk : int, optional
        Used in the group topk
    scaling_factor : float, optional
        Scaling factor for the routing scores
    score_function : str
        Currently support "softmax", "sigmoid" and "sqrtsoftplus".
    expert_bias : torch.Tensor, optional
        Could be used with the sigmoid/sqrtsoftplus score functions.

    Returns
    -------
    probs : torch.Tensor
        Routing probabilities in the same dtype as the "logits".
        Shape: [num_tokens, num_experts] (sparse, only topk entries are non-zero)
    routing_map : torch.Tensor
        Boolean mask indicating which experts were selected.
        Shape: [num_tokens, num_experts]
    """
    if logits.dtype == torch.float64:
        raise ValueError("Current TE does not support float64 router type.")
    
    num_tokens, num_experts = logits.shape
    
    # Precision notes:
    # - Logits are converted to fp32 for score functions.
    # - All the intermediate calculations are in fp32.
    # - The final probs are casted to the same dtype as the logits.
    if score_function == "softmax":
        if use_pre_softmax:
            scores = torch.softmax(logits, dim=-1, dtype=torch.float32)
            probs, top_indices = _compute_topk(scores, topk, num_groups, group_topk)
        else:
            scores, top_indices = _compute_topk(logits, topk, num_groups, group_topk)
            probs = torch.softmax(scores, dim=-1, dtype=torch.float32)
    elif score_function in ("sigmoid", "sqrtsoftplus"):
        if score_function == "sigmoid":
            scores = torch.sigmoid(logits.float())
        else:
            scores = torch.nn.functional.softplus(logits.float()).sqrt()
        if expert_bias is not None:
            scores_for_routing = scores + expert_bias.float()
            _, top_indices = _compute_topk(scores_for_routing, topk, num_groups, group_topk)
            scores = torch.gather(scores, dim=1, index=top_indices)
        else:
            scores, top_indices = _compute_topk(scores, topk, num_groups, group_topk)
        probs = scores / (scores.sum(dim=-1, keepdim=True) + 1e-20) if topk > 1 else scores
    else:
        raise ValueError(f"Invalid score_function: {score_function}")

    if scaling_factor:
        probs = probs * scaling_factor

    probs = probs.type_as(logits)

    # Build routing_map from top_indices
    if torch.are_deterministic_algorithms_enabled():
        # build [num_tokens, num_experts] from [num_tokens, topk]
        routing_probs = torch.zeros_like(logits)
        rows = torch.arange(num_tokens, device=logits.device).unsqueeze(1)
        routing_probs.index_put_((rows, top_indices), probs, accumulate=False)

        routing_map = torch.zeros_like(logits, dtype=logits.dtype)
        routing_map.index_put_(
            (rows, top_indices), torch.ones_like(probs, dtype=routing_map.dtype), accumulate=False
        )
        routing_map = routing_map.bool()
    else:
        # TODO Try using element-wise operations instead of scatter?
        routing_probs = torch.zeros_like(logits).scatter(1, top_indices, probs)
        routing_map = torch.zeros_like(logits).int().scatter(1, top_indices, 1).bool()

    return routing_probs, routing_map


# ===================== fused_compute_score_for_moe_aux_loss =====================


def fused_compute_score_for_moe_aux_loss(
    logits: torch.Tensor,
    topk: int,
    score_function: str,
) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    Fused compute scores for MoE aux loss for NPU.
    
    This is a placeholder implementation using small operators (non-fused).
    It provides the same interface as the TransformerEngine version but uses
    PyTorch operations instead of fused CUDA kernels.
    
    This implementation follows Megatron's approach and relies on PyTorch's
    automatic differentiation for gradient computation.
    
    Parameters
    ----------
    logits : torch.Tensor in fp32/bf16/fp16
        Router logits of shape [num_tokens, num_experts]
    topk : int
        Number of experts to select for each token
    score_function : str
        Currently support "softmax", "sigmoid" and "sqrtsoftplus".

    Returns
    -------
    routing_map : torch.Tensor
        Boolean mask indicating which experts were selected.
        Shape: [num_tokens, num_experts]
    scores : torch.Tensor
        Routing scores in fp32.
        Shape: [num_tokens, num_experts]
    """
    # Compute scores based on score_function
    if score_function == "softmax":
        scores = torch.softmax(logits, dim=-1, dtype=torch.float32)
    elif score_function == "sigmoid":
        # Cast logits to float32 before sigmoid for stability
        scores = torch.sigmoid(logits.to(torch.float32))
        scores = scores / (scores.sum(dim=-1, keepdim=True) + 1e-20)
    elif score_function == "sqrtsoftplus":
        # sqrtsoftplus: sqrt(softplus(x)) = sqrt(log(1 + exp(x)))
        scores = torch.nn.functional.softplus(logits.float()).sqrt().type_as(logits)
        scores = scores / (scores.sum(dim=-1, keepdim=True) + 1e-20)
    else:
        raise ValueError(f"Invalid score_function: {score_function}")

    # Get top-k indices
    _, top_indices = torch.topk(scores, k=topk, dim=1)
    
    # Build routing_map from top_indices
    routing_map = torch.zeros_like(logits).int().scatter(1, top_indices, 1).bool()

    return routing_map, scores


# ===================== fused_moe_aux_loss =====================


def fused_moe_aux_loss(
    probs: torch.Tensor,
    tokens_per_expert: torch.Tensor,
    total_num_tokens: int,
    num_experts: int,
    topk: int,
    coeff: float,
) -> torch.Tensor:
    """
    Fused MoE aux loss for NPU.
    
    This is a placeholder implementation using small operators (non-fused).
    It provides the same interface as the TransformerEngine version but uses
    PyTorch operations instead of fused CUDA kernels.
    
    This implementation follows Megatron's approach and relies on PyTorch's
    automatic differentiation for gradient computation.
    
    Parameters
    ----------
    probs : torch.Tensor in fp32/bf16/fp16
        Routing probabilities of shape [num_tokens, num_experts]
    tokens_per_expert : torch.Tensor in int32/int64/fp32/bf16
        The number of tokens per expert. Shape: [num_experts]
    total_num_tokens : int
        The total number of tokens used in the aux loss calculation.
    num_experts : int
        Number of experts
    topk : int
        Number of experts selected per token
    coeff : float
        The coefficient of the aux loss.

    Returns
    -------
    aux_loss : torch.Tensor
        A scalar tensor in the same dtype as the "probs".
    """
    # Compute auxiliary loss
    # The formula for the auxiliary loss is:
    #     loss = E * Σ_{i=1}^{E} (f_i * P_i)
    # where:
    #     f_i = 1 / (T * topk) * Σ_{x∈B} routing_map(x, i)
    #          (fraction of tokens dispatched to expert i)
    #     P_i = 1 / T * Σ_{x∈B} probs(x, i)
    #          (averaged router probability allocated for expert i)
    #     E is the number of experts
    #     T is the total number of tokens in the batch B
    
    aggregated_probs_per_expert = probs.sum(dim=0)
    aux_loss = torch.sum(aggregated_probs_per_expert * tokens_per_expert) * (
        num_experts * coeff / (topk * total_num_tokens * total_num_tokens)
    )
    
    return aux_loss