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

from collections.abc import Callable

import torch

try:
    from mt_dao.models.cvmm_triton_kernel import (
        cvmm_prepare_sel2,
    )
    from mt_dao.models.cvmm_triton_kernel_fp8_activations import (
        cvmm_fp8,
    )
except (ImportError, ModuleNotFoundError, RuntimeError):

    cvmm_fp8 = None  # type: ignore[assignment]
    cvmm_prepare_sel2 = None  # type: ignore[assignment]


def forward_cvmm_fp8_kernel(  # noqa: PLR0913, 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:
    """Forward pass of the SigmaMoE layer.

    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
        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
        The output tensor after applying the SigmaMoE layer, with shape
        (batch_size, seq_length, v_dim).

    Raises
    ------
    RuntimeError
        If the CVMM kernels are not available when selected.

    """
    if cvmm_fp8 is None or cvmm_prepare_sel2 is None:
        msg = "CVMM kernels are not available"
        raise RuntimeError(msg)
    # 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
    # Preprocess the selection indices. They will be needed for both layers and save
    # some time
    # NOTE: `sel_indices.reduction_weight` will be set to None at this stage
    sel_indices = cvmm_prepare_sel2(sel_index.int())

    # Up-projection layer for each head
    scores = cvmm_fp8(x, sel_indices, up_proj.view(n_experts, d_model, expert_size))
    # Apply activation function
    scores = activation(scores)

    # Down projection layer for each head
    sel_indices = sel_indices.clone()
    sel_indices.reduction_weight = sel_val
    assert sel_indices.out_index is not None, "Output index should not be None"
    sel_indices.sel_index = sel_indices.out_index
    sel_indices.out_index = None
    out = cvmm_fp8(scores, sel_indices, down_proj)

    return out.view(*x.shape[:-1], v_dim)
