"""Module for building SigmaMoE layers."""

from collections.abc import Callable
from typing import Any

import torch
from llmfoundry.models.layers.ffn import resolve_ffn_act_fn

from mt_dao.models.deptoe.mup_config import MuPConfig
from mt_dao.models.ffns.sigma_moe import SigmaMoE, SigmaMoEConfig


def build_sigma_moe(  # noqa: PLR0913, PLR0917
    d_model: int,
    ff_n_experts: int,
    n_repeats: int,
    ff_expert_size: int,
    ff_n_active_experts: int,
    v_dim: int | None = None,
    mup_config: dict | None = None,
    ffn_act_fn: dict | None = None,
    expert_dropout: float = 0.0,
    init_std: float = 0.02,
    r_std_multiplier: float = 1.0,
    e_std_multiplier: float = 1.0,
    *,
    expansion_ratio: float | None = None,  # noqa: ARG001
    device: str | None = None,  # noqa: ARG001
    bias: bool = False,  # noqa: ARG001
    fc_type: str = "torch",  # noqa: ARG001
    cvmm_kernel: bool = False,
    deepgemm_kernel: bool = False,
    h100_fp8_train: bool = False,
) -> torch.nn.Module:
    """Build a SigmaMoE layer with the given configuration.

    Parameters
    ----------
    d_model : int
        The dimension of the model.
    ff_n_experts : int
        The number of experts in the mixture of experts.
    n_repeats : int
        The number of repeats for the Universal Transformer (UT) mechanism.
    ff_expert_size : int
        The size of each expert, i.e., the dimensionality of the up-projection of
        each expert.
    ff_n_active_experts : int
        The number of active experts to select for each input. This is the number of
        experts that will be selected for each input during the forward pass.
    v_dim : int | None, optional
        The dimension of the value vectors (default is None, which uses d_model).
        It is the dimensionality of the down-projection of each expert.
    mup_config : dict | None, optional
        Configuration for the MuP mechanism. If None, defaults to an empty MuPConfig.
    ffn_act_fn : dict | None, optional
        The activation function to apply to the scores for selecting the experts.
        If None, defaults to F.silu.
    expert_dropout : float, optional
        The dropout rate for the experts (default is 0.0, meaning no dropout).
    init_std : float, optional
        The standard deviation for the initialization of the parameters
        (default is 0.02).
    r_std_multiplier : float, optional
        The multiplier for the standard deviation of the router (default is 1.0).
    e_std_multiplier : float, optional
        The multiplier for the standard deviation of the experts (default is 1.0).
    expansion_ratio : float | None, optional
        The expansion ratio for the experts. This is not used in the
        SigmaMoE layer, needed for compatibility with other FFNs.
    device : str | None, optional
        The device to place the model on (default is None,
        which uses the current device).
    bias : bool, optional
        Whether to use bias in the linear layers (default is False). This is not used in
        the SigmaMoE layer, needed for compatibility with other FFNs.
    fc_type : str, optional
        The type of fully connected layer to use (default is "torch"). This is not used
        in the SigmaMoE layer, needed for compatibility with other FFNs.
    cvmm_kernel : bool, optional
        Whether to use a custom kernel for the forward pass (default is False).
        If True, the forward pass will use a custom kernel for efficient computation.
    deepgemm_kernel: bool, optional
        If True, uses the DeepGEMM kernel for the feed-forward network. Default is
        False.
    h100_fp8_train: bool, optional
        If True, enables H100 FP8 training for the feed-forward network. Default is
        False.

    Returns
    -------
    nn.Module
        An instance of the SigmaMoE layer with the specified configuration.

    """
    return SigmaMoE(
        config=SigmaMoEConfig(
            d_model=d_model,
            ff_n_experts=ff_n_experts,
            n_repeats=n_repeats,
            n_blocks=1,  # Single block for now
            ff_expert_size=ff_expert_size,
            ff_n_active_experts=ff_n_active_experts,
            mup_config=MuPConfig(**mup_config) if mup_config else MuPConfig(),
            v_dim=v_dim,
            activation=(
                resolve_ffn_act_fn(ffn_act_fn)
                if ffn_act_fn
                else torch.nn.functional.silu
            ),
            expert_dropout=expert_dropout,
            cvmm_kernel=cvmm_kernel,
            deepgemm_kernel=deepgemm_kernel,
            h100_fp8_train=h100_fp8_train,
            init_std=init_std,
            r_std_multiplier=r_std_multiplier,
            e_std_multiplier=e_std_multiplier,
        ),
    )


# NOTE: we need this to avoid getting shouted at by mosaic
def sigma_moe_init_sp(
    module: torch.nn.Module,
    init_fn_: Callable,
    **kwargs: Any,  # noqa: ANN401
) -> bool:
    """Initialize the SigmaMoE module parameters.

    This function is used to initialize the parameters of the SigmaMoE module
    with the provided initialization function. It initializes the up_projection,
    down_projection, and expert_sel parameters of the SigmaMoE module using a
    specified function.

    Parameters
    ----------
    module : nn.Module
        The SigmaMoE module to initialize.
    init_fn_ : Callable
        The initialization function to use for the parameters.
    **kwargs : Any
        Additional keyword arguments (not used, just to capture any extra args).

    Returns
    -------
    bool
        Returns True if the initialization was successful, otherwise raises an
        assertion error if the module is not an instance of SigmaMoE.


    """
    del kwargs  # unused, just to capture any extra args

    if isinstance(module, SigmaMoE):
        init_fn_(module.up_proj)
        init_fn_(module.down_proj)
        if hasattr(module, "expert_sel"):
            init_fn_(module.expert_sel)

        return True

    return False
