"""SigmaMoE layer for mixture of experts with selection mechanism.

This module implements the SigmaMoE layer, which is a mixture of experts (MoE)
layer with a selection mechanism. It allows for dynamic selection of experts based on
the input data, enabling efficient computation and flexibility in model architecture.

Inspired from: https://github.com/RobertCsordas/moeut
Research Papers:
 - Approximating Two-Layer Feedforward Networks for Efficient Transformers,
   https://arxiv.org/abs/2310.10837

"""

# ruff: noqa: I001
from __future__ import annotations

import math
from dataclasses import dataclass
from typing import TYPE_CHECKING

import torch

from mt_dao.models.ffns.utils import entropy_reg

from mt_dao.models.ffns.forward_router import forward_router

from mt_dao.models.ffns.forward_no_kernel import (
    forward_dense_single_expert,
    forward_dense_all_experts,
    forward_sparse,
)


from mt_dao.models.ffns.forward_cvmm import forward_cvmm_kernel
from mt_dao.models.ffns.forward_gemm import forward_deepgemm

try:
    from mt_dao.models.ffns.forward_cvmm_h100_fp8 import (
        forward_cvmm_fp8_kernel,
    )
except (ImportError, ModuleNotFoundError):  # pragma: no cover - optional path
    forward_cvmm_fp8_kernel = None  # type: ignore[assignment]

if TYPE_CHECKING:
    from collections.abc import Callable
    from mt_dao.models.deptoe.mup_config import MuPConfig

try:
    from mt_dao.models.ffns.deepgemm_moe import (
        is_deepgemm_available,
    )
except (ImportError, ModuleNotFoundError):

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

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

        """
        return False


try:
    from mt_dao.models.cvmm_triton_kernel import (
        is_cvmm_available,
    )
except (ImportError, ModuleNotFoundError, RuntimeError):

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

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

        """
        return False


LIST_EXPERTS_PARAMS = [
    "ffn.up_proj",
    "ffn.down_proj",
    "ffn.expert_sel",
]


@dataclass
class SigmaMoEConfig:
    """Configuration for the SigmaMoE layer.

    Parameters
    ----------
    d_model : int
        Model hidden size.
    ff_n_experts : int
        Total number of experts.
    n_repeats : int
        The number of repeats for the Universal Transformer (UT) mechanism.
    n_blocks : int
        The number of blocks containing a SigmaMoE layer in the UT.
    ff_expert_size : int
        Size of each expert's up-projection dimension (S).
    ff_n_active_experts : int
        Number of active experts selected per token (top-k).
    mup_config : MuPConfig
        MuP configuration.
    activation : Callable[[torch.Tensor], torch.Tensor]
        Activation function applied after the up-projection.
    v_dim : int | None
        Output dimension V (defaults to d_model when None).
    expert_dropout : float
        Dropout applied to router scores during training.
    cvmm_kernel : bool
        Whether to use a custom Triton kernel path.
    deepgemm_kernel : bool
        Whether to use the DeepGEMM kernel.
    h100_fp8_train : bool
        Whether to enable H100 FP8 training.
    init_std : float
        Parameter init stddev baseline.
    r_std_multiplier : float
        Router stddev multiplier.
    e_std_multiplier : float
        Experts stddev multiplier.

    """

    d_model: int
    ff_n_experts: int
    n_repeats: int
    n_blocks: int
    ff_expert_size: int
    ff_n_active_experts: int
    mup_config: MuPConfig
    activation: Callable[[torch.Tensor], torch.Tensor] = torch.nn.functional.silu
    v_dim: int | None = None
    expert_dropout: float = 0.0
    cvmm_kernel: bool = False
    deepgemm_kernel: bool = False
    # Enable H100 FP8 compute during training by accelerating the up-projection
    # via DeepGEMM FP8 while keeping backward using the existing CVMM kernels.
    # Requires SM90+ and DeepGEMM installed. Falls back automatically otherwise.
    h100_fp8_train: bool = False
    init_std: float = 0.02
    r_std_multiplier: float = 1.0
    e_std_multiplier: float = 1.0


class SigmaMoE(torch.nn.Module):
    """SigmaMoE layer implements a mixture of experts with a selection mechanism."""

    def __init__(
        self,
        config: SigmaMoEConfig,
    ) -> None:
        """Initialize the SigmaMoE layer.

        Parameters
        ----------
        config : SigmaMoEConfig
            Configuration object containing parameters for the SigmaMoE layer.

        """
        super().__init__()
        self.config = config

        # NOTE: Adding a multiplier to the std of the initialization distribution for
        # the experts and the router
        self.stddev_experts_multiplier = config.e_std_multiplier
        self.stddev_router_multiplier = config.r_std_multiplier

        self.d_model = config.d_model
        self.v_dim = config.v_dim if config.v_dim is not None else config.d_model
        self.n_experts = config.ff_n_experts
        self.expert_size = config.ff_expert_size
        self.n_active_experts = config.ff_n_active_experts
        self.activation = config.activation
        self.expert_dropout = config.expert_dropout

        self.n_repeats = config.n_repeats
        self.n_blocks = config.n_blocks

        self.sel_hist: list[torch.Tensor] = []

        # NOTE: Deepseek does not use biases
        self.up_proj = torch.nn.Parameter(
            torch.empty(self.n_experts, self.expert_size, self.d_model),
        )
        self.down_proj = torch.nn.Parameter(
            torch.empty(self.n_experts, self.expert_size, self.v_dim),
        )
        if self.n_experts > 1:
            self.expert_sel = torch.nn.Parameter(
                torch.empty(self.n_experts, self.d_model),
            )

        self.use_cvmm_kernel = is_cvmm_available() and config.cvmm_kernel

        self.use_deepgemm = is_deepgemm_available() and config.deepgemm_kernel

        self.use_fp8_path = (
            self.config.h100_fp8_train
            and is_deepgemm_available()
            and is_cvmm_available()
        )

    def extra_repr(self) -> str:
        """Return a string representation of the SigmaMoE layer's parameters.

        The extra representation includes the shapes of the parameters `keys`,
        `values`, `expert_sel`.

        Returns
        -------
        str
            A string representation of the SigmaMoE layer's parameters.

        """
        extra_repr_str = ""
        if self.n_experts > 1:
            extra_repr_str += f"(expert_sel): Parameters({self.expert_sel.shape})\n"

        extra_repr_str += f"\n(self.up_projection): Parameters({self.up_proj.shape})\n"
        extra_repr_str += (
            f"\n(self.down_projection): Parameters({self.down_proj.shape})"
        )
        return extra_repr_str

    def _reset_parameters_sp(self) -> None:
        """Reset the parameters of the SigmaMoE layer using standard parametrization."""
        with torch.no_grad():
            scale = math.sqrt(2 / (self.n_repeats * self.n_blocks))

            if self.n_experts > 1:
                torch.nn.init.normal_(
                    self.expert_sel,
                    0,
                    scale / math.sqrt(self.d_model),
                )
                self.renorm_keep_std(self.expert_sel, dim=1)

            torch.nn.init.normal_(
                self.up_proj,
                0,
                scale / math.sqrt(self.d_model),
            )
            torch.nn.init.normal_(
                self.down_proj,
                0,
                scale / math.sqrt(self.n_experts * self.expert_size),
            )

    def _reset_parameters_mup(self) -> None:
        """Reset the parameters of the SigmaMoE layer using MuP."""
        with torch.no_grad():
            if self.n_experts > 1:
                torch.nn.init.normal_(
                    self.expert_sel,
                    0,
                    (
                        self.config.init_std
                        * self.stddev_router_multiplier
                        / math.sqrt(self.config.mup_config.mup_width_multiplier)
                        if self.config.mup_config.mup_scale_expert_sel
                        else (
                            math.sqrt(2 / (self.n_repeats * self.n_blocks))
                            / math.sqrt(self.d_model)
                        )
                    ),
                )
                self.renorm_keep_std(self.expert_sel, dim=1)

            torch.nn.init.normal_(
                self.up_proj,
                0,
                self.config.init_std
                / math.sqrt(self.config.mup_config.mup_width_multiplier),
            )
            # Use boolean to control weather we have n_experts in denominator of muP
            # NOTE: Maybe ablate how the number of experts relates to width
            torch.nn.init.normal_(
                self.down_proj,
                0,
                self.config.init_std
                * self.stddev_experts_multiplier
                / (
                    math.sqrt(
                        self.config.mup_config.mup_width_multiplier,
                    )
                ),
            )

    def reset_parameters(self) -> None:
        """Reset the parameters of the SigmaMoE layer using SP or MuP."""
        if self.config.mup_config.mup_enabled:
            # Reset parameters for MuP
            self._reset_parameters_mup()
        else:
            self._reset_parameters_sp()

    @staticmethod
    def renorm_keep_std(weight: torch.Tensor, dim: int = 0) -> None:
        """Renormalize the weight tensor to have unit norm on the specified dimension.

        This function ensures that the standard deviation of the weights remains
        unchanged after normalization.

        Parameters
        ----------
        weight : torch.Tensor
            The weight tensor to be renormalized.
        dim : int, optional
            The dimension along which to normalize the weights (default is 0).

        """
        with torch.no_grad():
            std = weight.std()
            weight.div_(weight.norm(dim=dim, keepdim=True))
            weight.mul_(std / weight.std())

    def get_reg_loss(self) -> torch.Tensor:
        """Calculate the regularization loss based on the selection history.

        Returns
        -------
        torch.Tensor
            The regularization loss computed from the selection history.

        """
        if not self.sel_hist:
            return torch.Tensor(0)

        # NOTE: The selection history is a list of 't' tensors, where 't' is the number
        # of (UT) repetitions executed of the group, each
        # with shape (batch_size, seq_length, n_experts)
        # Average over time and layers.
        loss = entropy_reg(torch.stack(self.sel_hist, dim=-2).flatten(-3, -2), -2)
        self.sel_hist = []
        return loss

    def forward(
        self,
        x: torch.Tensor,
    ) -> torch.Tensor:
        """Forward pass of the SigmaMoE layer.

        At the beginning of the forward, the selection history is cleared. As such, it
        is kept only for the duration of one batch step.

        Parameters
        ----------
        x : torch.Tensor
            The main data path tensor, expected shape is
            (batch_size, seq_length, d_model).

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

        """
        # Clear the selection history
        self.sel_hist = []

        # Operate the correct forward pass based on the configuration

        # NOTE: If n_repeats == 1, we default to a pre-layernorm transformer
        # Which we assume comes from the block
        # the original MoEUT paper does not normalize the values
        # in order to allow better scaling with model depth
        # for use CompleteP fulfills the same purpose
        # the SigmaMOE paper also uses a pre-layernorm transformer
        if self.n_active_experts == self.n_experts == 1:
            return forward_dense_single_expert(
                x=x,
                up_proj=self.up_proj,
                down_proj=self.down_proj,
                activation=self.activation,
            )
        # Obtain the selection scores and indices from the router
        all_scores, topk_sel_val, topk_sel_index = forward_router(
            x=x,
            expert_sel=self.expert_sel,
            n_active_experts=self.n_active_experts,
            sel_input=x,
            expert_dropout=self.expert_dropout,
            is_training=self.training,
        )
        # Append the selection scores to the history
        self.sel_hist.append(all_scores)

        # Prefer DeepGEMM path when available; fallback to vectorized torch
        if not self.training and self.use_deepgemm and self.n_active_experts > 1:
            return forward_deepgemm(
                x=x,
                sel_val=topk_sel_val,
                sel_index=topk_sel_index,
                up_proj=self.up_proj,
                down_proj=self.down_proj,
                activation=self.activation,
            )

        # Training-time FP8 forward path on H100 using DeepGEMM + CVMM backward
        if (
            self.training
            and self.use_fp8_path
            and forward_cvmm_fp8_kernel is not None
            and self.n_active_experts > 1
        ):
            topk_sel_index = (
                (
                    torch.arange(
                        self.n_experts, device=topk_sel_val.device, dtype=torch.int64,
                    )
                    .view(1, 1, -1)
                    .expand(x.shape[0], x.shape[1], -1)
                )
                if topk_sel_index is None
                else topk_sel_index
            )
            return forward_cvmm_fp8_kernel(
                x=x,
                sel_val=topk_sel_val,
                sel_index=topk_sel_index,
                up_proj=self.up_proj,
                down_proj=self.down_proj,
                activation=self.activation,
            )
        if self.use_cvmm_kernel:
            topk_sel_index = (
                (
                    torch.arange(
                        self.n_experts, device=topk_sel_val.device, dtype=torch.int64,
                    )
                    .view(1, 1, -1)
                    .expand(x.shape[0], x.shape[1], -1)
                )
                if topk_sel_index is None
                else topk_sel_index
            )
            return forward_cvmm_kernel(
                x=x,
                sel_val=topk_sel_val,
                sel_index=topk_sel_index,
                up_proj=self.up_proj,
                down_proj=self.down_proj,
                activation=self.activation,
            )

        # Fast dense path when all experts are active (no indices returned)
        if topk_sel_index is None:
            return forward_dense_all_experts(
                x=x,
                sel_val=topk_sel_val,
                up_proj=self.up_proj,
                down_proj=self.down_proj,
                activation=self.activation,
            )
        return forward_sparse(
            x=x,
            sel_val=topk_sel_val,
            sel_index=topk_sel_index,
            up_proj=self.up_proj,
            down_proj=self.down_proj,
            activation=self.activation,
        )
