"""Module implementing the forward pass of a SigmaMoE layer using DeepGEMM kernels."""

from collections.abc import Callable

import torch

try:
    from mt_dao.models.ffns.deepgemm_moe import (
        dense_forward_deepgemm as _dense_forward_deepgemm,
    )
    from mt_dao.models.ffns.deepgemm_moe import (
        moe_forward_deepgemm as _moe_forward_deepgemm,
    )
except (ImportError, ModuleNotFoundError):

    def is_deepgemm_available() -> bool:
        """Return False when DeepGEMM helper is not importable.

        Returns
        -------
        bool
            Always False.

        """
        return False

    _dense_forward_deepgemm = None  # type: ignore[assignment]
    _moe_forward_deepgemm = None  # type: ignore[assignment]


def forward_deepgemm(  # noqa: PLR0913, PLR0917
    x: torch.Tensor,
    sel_val: torch.Tensor,
    sel_index: torch.Tensor | None,
    up_proj: torch.Tensor,
    down_proj: torch.Tensor,
    activation: Callable[[torch.Tensor], torch.Tensor],
) -> torch.Tensor:
    """Forward pass using DeepGEMM grouped FP8 kernels.

    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). This will be None if all experts are
        selected.
    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 with shape (batch_size, seq_length, v_dim).

    Raises
    ------
    RuntimeError
        If the DeepGEMM helper is not available when selected.

    """
    # Delegate to the DeepGEMM helper
    if _moe_forward_deepgemm is None or _dense_forward_deepgemm is None:
        msg = "DeepGEMM helper not available"
        raise RuntimeError(msg)
    if sel_index is None:
        return _dense_forward_deepgemm(
        x=x,
        sel_weight=sel_val,
        up_proj=up_proj,
        down_proj=down_proj,
        activation=activation,
    )
    return _moe_forward_deepgemm(
        x=x,
        sel_index=sel_index,
        sel_weight=sel_val,
        up_proj=up_proj,
        down_proj=down_proj,
        activation=activation,
    )
