"""Module for implementing the forward of an MoE router."""

import torch


def forward_router(  # noqa: PLR0913
    x: torch.Tensor,
    expert_sel: torch.nn.Parameter,
    n_active_experts: int,
    sel_input: torch.Tensor | None = None,
    expert_dropout: float = 0.0,
    *,
    is_training: bool = True,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
    """Compute router scores and return top-k weights and indices.

    Parameters
    ----------
    x : torch.Tensor
        Input tensor of shape (batch, seq, d_model).
    expert_sel : torch.nn.Parameter
        Expert selection weights of shape (d_model, n_experts).
    n_active_experts : int
        Number of active experts to use for topk selection.
    sel_input : torch.Tensor | None, optional
        Optional tensor to use for selection scoring; defaults to x.
    expert_dropout : float, optional
        Dropout rate for expert selection; defaults to 0.0.
    is_training : bool, optional
        Whether the model is in training mode; defaults to True.

    Returns
    -------
    tuple[torch.Tensor, torch.Tensor | None]
        sel_val and sel_index, both of shape (batch, seq, n_active_experts).
        sel_index is None if n_active_experts == n_experts.

    """
    assert n_active_experts > 1, (
        "This method is only for the case where the number of active experts is "
        "greater than 1 (equal to the number of experts)."
    )
    # Get the number of experts from the expert selector
    n_experts = expert_sel.shape[-1]

    # Selection score calculation with shape (batch_size, seq_length, n_experts)
    sel = torch.nn.functional.linear(
        sel_input if sel_input is not None else x,
        expert_sel,
        bias=None,
    )

    # Selection activation and topk
    sel = torch.nn.functional.sigmoid(sel)

    with torch.no_grad():
        # Apply dropout to the selection scores if required
        if is_training and expert_dropout > 0:
            mask = torch.rand_like(sel) < expert_dropout
            sel = sel.masked_fill(mask, float("-inf"))

    # If all experts are active, avoid topk overhead: select all experts per token
    if n_active_experts == n_experts:
        return sel, sel, None
    # Select the top-k elements from the selection scores
    topk_sel_val, topk_sel_index = sel.topk(n_active_experts, dim=-1, sorted=False)
    return sel, topk_sel_val, topk_sel_index
