"""Module for implementing the forward pass of a dense layer in a SigmaMoE model."""

from collections.abc import Callable

import torch


def forward_dense_single_expert(
    x: torch.Tensor,
    up_proj: torch.nn.Parameter,
    down_proj: torch.nn.Parameter,
    activation: Callable[[torch.Tensor], torch.Tensor],
) -> torch.Tensor:
    """Forward pass of the SigmaMoE layer with simplified logic.

    This method assumes that all experts are active and does not use the selection
    mechanism. It is intended for cases where all experts are used, such as in
    training or evaluation scenarios where the selection mechanism is not needed.

    Parameters
    ----------
    x : torch.Tensor
        The input tensor to the SigmaMoE layer, expected shape is
        (batch_size, seq_length, d_model). This is usually the main (residual)
        data path summed with the output of the multi-head attention scores.
    up_proj : torch.nn.Parameter
        The up-projection weights for the experts.
    down_proj : torch.nn.Parameter
        The down-projection weights for the experts.
    activation : Callable[[torch.Tensor], torch.Tensor]
        The activation function to apply after the up-projection.

    Returns
    -------
    torch.Tensor
        The output tensor after applying the SigmaMoE layer, with shape
        (batch_size, seq_length, v_dim).

    """
    # Get the dimensions from the input tensor
    _batch_size, _seq_length, d_model = x.shape
    # Get the experts dimensions from the up projection and down projection
    expert_size = up_proj.view(-1, d_model).shape[0]
    v_dim = down_proj.view(-1, expert_size).shape[0]

    # "Up-projection" layer for each head
    x_up_proj = activation(
        torch.nn.functional.linear(
            x,
            up_proj.view(-1, d_model),
            None,
        ),
    )

    # Down projection layer for each head
    x_down_proj = torch.nn.functional.linear(
        x_up_proj,
        down_proj.view(expert_size, v_dim).t(),
        None,
    )

    return x_down_proj.view(*x.shape)


def forward_dense_all_experts(
    x: torch.Tensor,
    sel_val: torch.Tensor,
    up_proj: torch.Tensor,
    down_proj: torch.Tensor,
    activation: Callable[[torch.Tensor], torch.Tensor],
) -> torch.Tensor:
    """Dense fast path when all experts are active (no indices provided).

    This path avoids any Python loops or scatter/gather by flattening experts
    and using two fused linear operations with reshapes. It computes, for each
    token, the sum over experts of the up-projection (with activation), scaled
    by the router weights, followed by the expert-specific down-projections.

    Shapes
    ------
    x: [B, T, D]
    sel_val: [B, T, E]
    up_proj: [E, S, D]
    down_proj: [E, S, V]

    Returns
    -------
    torch.Tensor
        The output tensor with shape [B, T, V].

    """
    # Dimensions
    bsz, seqlen, d_model = x.shape
    n_experts, expert_size, _ = up_proj.shape
    v_dim = down_proj.shape[-1]

    # anonty checks (lightweight)
    assert up_proj.shape[-1] == d_model, "up_proj last dim must match x hidden size"
    assert down_proj.shape[0] == n_experts, "down_proj E must match up_proj E"
    assert down_proj.shape[1] == expert_size, "down_proj S must match up_proj S"

    # Flatten tokens
    x_flat = x.view(bsz * seqlen, d_model)

    # Up-projection over all experts at once: [N, D] -> [N, E*S]
    # F.linear: out = x @ W^T, with W:[E*S, D]
    u = torch.nn.functional.linear(
        x_flat,
        up_proj.view(n_experts * expert_size, d_model),
        bias=None,
    )
    u = activation(u)

    # Apply router gates per expert (broadcast to expert_size entries).
    # Router may have applied dropout by setting some gates to -inf; map those to 0.
    gates = sel_val.view(bsz * seqlen, n_experts)
    gates = torch.where(torch.isfinite(gates), gates, torch.zeros_like(gates))
    gates = (
        gates.unsqueeze(-1)
        .expand(-1, -1, expert_size)
        .view(bsz * seqlen, n_experts * expert_size)
    )
    u.mul_(gates)

    # Down-projection and sum across experts implicitly via dense matmul
    # W_down_combined:[E*S, V] -> use F.linear with weight:[V, E*S]
    y_flat = torch.nn.functional.linear(
        u,
        down_proj.view(n_experts * expert_size, v_dim).t(),
        bias=None,
    )  # [N, V]

    return y_flat.view(bsz, seqlen, v_dim)


def forward_sparse(  # noqa: PLR0913, PLR0914, PLR0917
    x: torch.Tensor,
    sel_val: torch.Tensor,
    sel_index: torch.Tensor,
    up_proj: torch.nn.Parameter,
    down_proj: torch.nn.Parameter,
    activation: Callable[[torch.Tensor], torch.Tensor],
) -> torch.Tensor:
    """Sparse grouped top-k fallback avoiding dense [B, T, E] tensors.

    Groups selected tokens by expert and runs per-expert dense matmuls:
    U_e = X_e @ W_up[e]^T; apply activation in-place; pre-scale by gate;
    Y_e = U_e @ W_down[e]; scatter-add back to output.

    Parameters
    ----------
    x : torch.Tensor
        The input tensor to the SigmaMoE layer, expected shape is
        (batch_size, seq_length, d_model).
    sel_val : torch.Tensor
        The selection values tensor, expected shape is
        (batch_size, seq_length, n_active_experts).
    sel_index : torch.Tensor | None
        The selection indices tensor, expected shape is
        (batch_size, seq_length, n_active_experts).
    up_proj : torch.nn.Parameter
        The up projection weights, expected shape is
        (n_experts, d_model, expert_size).
    down_proj : torch.nn.Parameter
        The down projection weights, expected shape is
        (n_experts, expert_size, v_dim).
    activation : Callable[[torch.Tensor], torch.Tensor]
        The activation function to apply to the output.

    Returns
    -------
    torch.Tensor
        Output tensor of shape [B, T, V].

    """
    # Get the dimensions from the input tensor
    batch_size, seq_length, d_model = x.shape
    # Get the experts dimensions from the up projection and down projection
    n_experts, _expert_size, d_model = up_proj.shape
    _n_experts, _expert_size, v_dim = down_proj.shape

    # Flatten tokens
    x_flat = x.reshape(-1, d_model)  # [N, D]
    n_tokens = batch_size * seq_length
    n_active_experts = sel_index.shape[-1]

    # Build (row_token, row_expert, row_weight) for selected experts only
    row_token = torch.arange(n_tokens, device=x.device).repeat_interleave(
        n_active_experts,
    )
    row_expert = sel_index.reshape(-1)
    row_weight = sel_val.reshape(-1)

    # Group by expert id
    order = torch.argsort(row_expert)
    row_token = row_token[order]
    row_expert = row_expert[order]
    row_weight = row_weight[order]

    counts = torch.bincount(row_expert, minlength=n_experts)
    offsets = torch.nn.functional.pad(counts.cumsum(0), (1, 0))[:-1]

    y_flat = x.new_zeros((n_tokens, v_dim))
    # Process each expert slice
    for e in range(n_experts):
        cnt = int(counts[e].item())
        if cnt == 0:
            continue
        start = int(offsets[e].item())
        sl = slice(start, start + cnt)
        tok_idx = row_token[sl]
        gate = row_weight[sl]
        x_e = x_flat.index_select(0, tok_idx)  # [cnt, D]

        # Up matmul per expert
        u = torch.matmul(x_e, up_proj[e].transpose(0, 1))
        u = activation(u)
        # Pre-scale by gate
        u.mul_(gate.unsqueeze(1))

        # Down matmul per expert
        y_e = torch.matmul(u, down_proj[e])

        # Scatter-add into output
        y_flat.index_add_(0, tok_idx, y_e)

    return y_flat.view(batch_size, seq_length, v_dim)
