"""Implementation of the SwitchHead layer for attention MoE with selection mechanism.

Inspired by the original implementation of SwitchHead in the MoeUT library: https://github.com/RobertCsordas/moeut

Research Papers:
 - MoEUT: Mixture-of-Experts Universal Transformers, https://arxiv.org/abs/2405.16039
 - SwitchHead: Accelerating Transformers with Mixture-of-Experts Attention, https://arxiv.org/abs/2312.07987

"""

import math
from collections import defaultdict
from dataclasses import dataclass
from logging import WARNING
from typing import cast

import torch
import torch.nn.functional as F  # noqa: N812
from flwr.common import log as flwr_log
from git import TYPE_CHECKING

from mt_dao.models.deptoe.mup_config import MuPConfig
from mt_dao.models.ffns.utils import entropy_reg
from mt_dao.models.rope.rope_encoding import RotaryPositionalEncoding
from mt_dao.models.switch_head.utils import AttentionMask, KVCache

LIST_EXPERTS_PARAMS = [
    "attn.v",
    "attn.o",
    "attn.sel_v",
    "attn.sel_o",
]

if TYPE_CHECKING:
    from mt_dao.models.cvmm_triton_kernel import CVMMSel


@dataclass
class SwitchHeadConfig:
    """Configuration class for SwitchHead.

    This class holds the configuration parameters for the SwitchHead layer,
    including state size, number of heads, number of experts, dropout rates, and
    other parameters related to the mixture of experts mechanism.

    Attributes
    ----------
    d_model : int
        The size of the input and output state.
    n_heads : int
        The number of attention heads.
    n_repeats : int
        The number of repeats for the Universal Transformer (UT) mechanism.
    n_blocks : int
        The number of blocks containing a SwitchHead layer in the UT.
    use_peri_ln : bool
        If True, uses a peripheral layer normalization (PeriLN) strategy from
        https://arxiv.org/abs/2502.02732 instead of the layer normalization strategy
        proposed in https://arxiv.org/abs/2405.16039
    attn_heads_n_experts : int
        The number of experts in the mixture of experts.
    attn_heads_dropout : float
        The dropout rate applied to the query vectors just before computing the
        attention scores.
    attn_heads_projection_size : int | None
        The size of the projection for each of the attention heads.
    attn_heads_expert_dropout : float
        The dropout rate on the experts selection scores.
    n_active_experts : int
        The number of SwitchHead active experts.
    cvmm_kernel : bool
        If True, uses a custom kernel for the attention mechanism. Default is False.
    init_std : float
        The standard deviation for the initialization of the parameters in the
        SwitchHead layer. Default is 0.02.
    no_bias : bool
        If True, disables the biases in the model for linear layers/layernorms.
    use_flash_attn : bool
        If True, uses FlashAttention for the attention mechanism. Default is True.
    mup_config : MuPConfig | None
        Configuration for the MuP (Maximal Update Parametrisation) mechanism.

    """

    d_model: int
    n_heads: int
    n_repeats: int
    n_blocks: int
    mup_config: MuPConfig
    attn_heads_n_experts: int
    attn_heads_dropout: float = 0.0
    attn_heads_projection_size: int | None = None
    attn_heads_expert_dropout: float = 0.0
    attn_heads_n_active_experts: int = 2
    cvmm_kernel: bool = False
    use_peri_ln: bool = False
    init_std: float = 0.02
    no_bias: bool = False
    use_flash_attn: bool = True


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

    def __init__(
        self,
        config: SwitchHeadConfig,
        rope_module: RotaryPositionalEncoding | None = None,
    ) -> None:
        """Initialize the SwitchHead layer.

        Parameters
        ----------
        config : SwitchHeadConfig
            The configuration object containing parameters for the SwitchHead layer.
        rope_module : RotaryPositionalEncoding | None, optional
            A module for applying rotary positional embeddings, by default None.

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

        self.input_size = config.d_model
        self.output_size = config.d_model
        self.attn_heads_expert_dropout = config.attn_heads_expert_dropout
        self.attn_heads_n_active_experts = config.attn_heads_n_active_experts
        self.attn_heads_n_experts = config.attn_heads_n_experts
        self.cvmm_kernel = config.cvmm_kernel
        self.use_flash_attn = config.use_flash_attn

        # Flag for using the peripheral layer normalization (PeriLN)
        self.use_peri_ln = config.use_peri_ln

        self.rope_module = rope_module

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

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

        self.n_heads = config.n_heads
        self.attn_heads_dropout = (
            torch.nn.Dropout(config.attn_heads_dropout)
            if config.attn_heads_dropout > 0
            else lambda x: x
        )
        self.attn_heads_projection_size = config.attn_heads_projection_size or (
            config.d_model // config.n_heads
        )
        # TODO(anon): Check why bias is missing here
        self.q = torch.nn.Linear(
            self.input_size,
            self.attn_heads_projection_size * self.n_heads,
            bias=False,
        )
        self.k = torch.nn.Linear(
            self.input_size,
            self.attn_heads_projection_size * self.n_heads,
            bias=False,
        )

        # NOTE: Deepseek does not use biases
        if self.attn_heads_n_experts > 1:
            self.v = torch.nn.Parameter(
                torch.empty(
                    self.attn_heads_n_experts,
                    self.n_heads,
                    self.input_size,
                    self.attn_heads_projection_size,
                ),
            )
            self.o = torch.nn.Parameter(
                torch.empty(
                    self.attn_heads_n_experts,
                    self.n_heads,
                    self.output_size,
                    self.attn_heads_projection_size,
                ),
            )
            self.sel_v = torch.nn.Parameter(
                torch.empty(
                    self.attn_heads_n_experts,
                    self.n_heads,
                    self.input_size,
                ),
            )
            self.sel_o = torch.nn.Parameter(
                torch.empty(
                    self.attn_heads_n_experts,
                    self.n_heads,
                    self.input_size,
                ),
            )

        else:
            self.v = torch.nn.Parameter(
                torch.empty(
                    self.n_heads * self.attn_heads_projection_size,
                    self.input_size,
                ),
            )
            self.o = torch.nn.Parameter(
                torch.empty(
                    self.output_size,
                    self.n_heads * self.attn_heads_projection_size,
                ),
            )

            self.sel_o = torch.nn.Parameter(
                torch.empty(
                    self.attn_heads_n_experts,
                    self.n_heads,
                    self.input_size,
                ),
            )

        if (
            self.config.mup_config.mup_enabled
            and not self.config.mup_config.mup_disable_attention_scaling
        ):
            # NOTE: the original MuP implementation
            # uses a scaling factors of 1/d_model rather than the conventional
            # 1/sqrt(d_model), this is to account for potential correlations
            # that may arise between the attention heads
            # this change does not seem to have a large impact
            # according to the NanoGPT muP benchmark results
            self.register_buffer(
                "scale",
                torch.full([1], 1.0 / self.attn_heads_projection_size),
                persistent=False,
            )
        else:
            self.register_buffer(
                "scale",
                torch.full([1], 1.0 / math.sqrt(self.attn_heads_projection_size)),
                persistent=False,
            )

        # Input LayerNorm
        self.input_ln = torch.nn.LayerNorm(self.input_size)
        # Add the output PeriLN layer
        # NOTE: With or without PeriLN we do have an input LayerNorm
        if self.use_peri_ln:
            self.peri_ln_out = torch.nn.LayerNorm(self.output_size)

        # Initialize the experts assignments counters
        self.experts_assignments_counter_v: dict[int, dict[int, int]] = defaultdict(
            lambda: defaultdict(int),
        )
        self.experts_assignments_counter_o: dict[int, dict[int, int]] = defaultdict(
            lambda: defaultdict(int),
        )
        # Initialize the experts scores statistics
        self.experts_scores_statistics_v: dict[int, dict[int, dict[str, float]]] = (
            defaultdict(
                lambda: defaultdict(lambda: defaultdict(float)),
            )
        )
        self.experts_scores_statistics_o: dict[int, dict[int, dict[str, float]]] = (
            defaultdict(
                lambda: defaultdict(lambda: defaultdict(float)),
            )
        )
        if self.use_flash_attn:
            from flash_attn import flash_attn_qkvpacked_func  # noqa: PLC0415

            self.flash_attn_qkvpacked_func = flash_attn_qkvpacked_func

        if self.config.cvmm_kernel:
            from mt_dao.models.cvmm_triton_kernel import (  # noqa: PLC0415
                cvmm,
                cvmm_prepare_sel2,
            )

            self.cvmm = cvmm
            self.cvmm_prepare_sel2 = cvmm_prepare_sel2

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

        The extra representation includes the shapes of the parameters `v`, `o`,
        `sel_v`, and `sel_o`.

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

        """
        extra_repr_str = ""
        if not (self.attn_heads_n_active_experts == self.attn_heads_n_experts == 1):
            extra_repr_str = f"(v): Parameters({self.v.shape})\n"

            extra_repr_str += f"(o): Parameters({self.o.shape})\n"
            if hasattr(self, "sel_v"):
                extra_repr_str += f"(sel_v): Parameters({self.sel_v.shape})\n"
            if hasattr(self, "sel_o"):
                extra_repr_str += f"(sel_o): Parameters({self.sel_o.shape})"
        return extra_repr_str

    def _reset_parameters_sp(self) -> None:
        """Reset the parameters of the SwitchHead layer using SP."""
        with torch.no_grad():
            scale = math.sqrt(2 / (self.n_repeats * self.n_blocks))
            if self.attn_heads_n_experts > 1:
                torch.nn.init.normal_(
                    self.sel_v,
                    0,
                    scale / math.sqrt(self.input_size),
                )
                self.renorm_rows(self.sel_v)

                torch.nn.init.normal_(self.sel_o, 0, scale / math.sqrt(self.input_size))
                self.renorm_rows(self.sel_o)

            torch.nn.init.normal_(
                self.k.weight,
                0,
                scale / math.sqrt(self.input_size),
            )
            torch.nn.init.normal_(
                self.q.weight,
                0,
                scale / math.sqrt(self.input_size),
            )
            torch.nn.init.normal_(
                self.v if isinstance(self.v, torch.nn.Parameter) else self.v.weight,
                0,
                scale / math.sqrt(self.input_size),
            )
            torch.nn.init.normal_(
                self.o if isinstance(self.o, torch.nn.Parameter) else self.o.weight,
                0,
                scale / math.sqrt(self.n_heads * self.attn_heads_projection_size),
            )

    def _reset_parameters_mup(self) -> None:
        """Reset the parameters of the SwitchHead layer using MuP."""
        with torch.no_grad():
            scale = (
                math.sqrt(2 / (self.n_repeats * self.n_blocks))
                / math.sqrt(self.input_size)
                if not self.config.mup_config.mup_scale_expert_sel
                else None
            )
            if self.attn_heads_n_experts > 1:
                torch.nn.init.normal_(
                    self.sel_v,
                    0,
                    self.config.init_std
                    / math.sqrt(self.config.mup_config.mup_width_multiplier)
                    if scale is None
                    else (scale / math.sqrt(self.input_size)),
                )
                self.renorm_rows(self.sel_v)

                torch.nn.init.normal_(
                    self.sel_o,
                    0,
                    self.config.init_std
                    / math.sqrt(self.config.mup_config.mup_width_multiplier)
                    if scale is None
                    else scale / math.sqrt(self.input_size),
                )
                self.renorm_rows(self.sel_o)

            torch.nn.init.normal_(
                self.k.weight,
                0,
                self.config.init_std
                / math.sqrt(self.config.mup_config.mup_width_multiplier),
            )
            torch.nn.init.normal_(
                self.q.weight,
                0,
                self.config.init_std
                / math.sqrt(self.config.mup_config.mup_width_multiplier),
            )
            torch.nn.init.normal_(
                self.v if isinstance(self.v, torch.nn.Parameter) else self.v.weight,
                0,
                self.config.init_std
                / math.sqrt(self.config.mup_config.mup_width_multiplier),
            )
            torch.nn.init.normal_(
                self.o if isinstance(self.o, torch.nn.Parameter) else self.o.weight,
                0,
                self.config.init_std
                / math.sqrt(self.config.mup_config.mup_width_multiplier),
            )
            flwr_log(
                WARNING,
                "Resetting SwitchHead parameters with MuP std: %s",
                self.config.init_std
                / math.sqrt(self.config.mup_config.mup_width_multiplier),
            )

    def reset_parameters(self) -> None:
        """Reset the parameters of the SwitchHead layer."""
        if hasattr(self, "config") and self.config.mup_config.mup_enabled:
            # Reset parameters for MuP
            self._reset_parameters_mup()
        else:
            # Reset parameters for SP
            self._reset_parameters_sp()

    def _reset_experts_assignments_counters(self) -> None:
        """Reset the experts assignments counter."""
        for attn_idx in self.experts_assignments_counter_v:
            for expert_idx in self.experts_assignments_counter_v[attn_idx]:
                self.experts_assignments_counter_v[attn_idx][expert_idx] = 0
        for attn_idx in self.experts_assignments_counter_o:
            for expert_idx in self.experts_assignments_counter_o[attn_idx]:
                self.experts_assignments_counter_o[attn_idx][expert_idx] = 0

    def get_experts_scores_statistics(
        self,
        experts_scores: torch.Tensor,
        *,
        is_v: bool = True,
    ) -> None:
        """Calculate statistics for the experts scores.

        The statistics include mean, standard deviation, median, maximum, and minimum
        for each expert's scores.

        Parameters
        ----------
        experts_scores : torch.Tensor
            A tensor containing the scores for each expert, expected shape is
            (n_heads, attn_heads_n_experts, batch_size * seq_length).
        is_v : bool, optional
            If True, use the statistics for the "v" experts, otherwise use the "o"
            experts.

        """
        experts_scores_statistics = (
            self.experts_scores_statistics_v
            if is_v
            else self.experts_scores_statistics_o
        )
        for attn_head_idx in range(self.n_heads):
            for expert_idx in range(self.attn_heads_n_experts):
                expert_scores = experts_scores[attn_head_idx][expert_idx]
                experts_scores_statistics[attn_head_idx][expert_idx]["mean"] = (
                    expert_scores.mean().item()
                )
                experts_scores_statistics[attn_head_idx][expert_idx]["std"] = (
                    expert_scores.std().item()
                )
                experts_scores_statistics[attn_head_idx][expert_idx]["median"] = (
                    expert_scores.median().item()
                )
                experts_scores_statistics[attn_head_idx][expert_idx]["max"] = (
                    expert_scores.max().item()
                )
                experts_scores_statistics[attn_head_idx][expert_idx]["min"] = (
                    expert_scores.min().item()
                )

    @staticmethod
    def renorm_rows(x: torch.Tensor) -> None:
        """Renormalize the rows of the tensor to have unit norm while keeping the std.

        Parameters
        ----------
        x : torch.Tensor
            The tensor to be renormalized. It is expected to have shape [..., N, C],
            where N is the number of rows and C is the number of columns.

        """
        with torch.no_grad():
            std_t = x.std(dim=-1, keepdim=True)
            x.div_(x.norm(dim=-1, keepdim=True))
            x.mul_(std_t / x.std())

    @staticmethod
    def get_mask_tensor(
        src_len: int,
        mask: AttentionMask | None,
    ) -> torch.Tensor | None:
        """Get the mask tensor for the attention mechanism.

        Parameters
        ----------
        src_len : int
            The length of the source sequence.
        mask : AttentionMask | None
            The attention mask containing position and source length masks.

        Returns
        -------
        torch.Tensor | None
            A tensor representing the mask for the attention mechanism, or None if no
            mask is provided.

        """
        # NOTE: The shape of the `mask.position_mask` is expected to be
        # [..., N_out, N_in], where N_out is the output length and N_in is the input
        # length. The shape of the `mask.src_length_mask` is expected to be
        # [B, ...., N_in], where B is the batch size and the ellipsis (...) represents
        # any number of additional dimensions. Their values are True where the
        # corresponding positions should be masked.
        if mask is None:
            return None
        position_mask = mask.position_mask
        src_length_mask = mask.src_length_mask
        if position_mask is None and src_length_mask is None:
            return None

        if position_mask is not None:
            n_pad = src_len - position_mask.shape[-1]
            if n_pad > 0:
                pm = F.pad(position_mask, (n_pad, 0), "constant", value=False)
            else:
                pm = position_mask
            return (
                pm
                if src_length_mask is None
                else pm | src_length_mask.unsqueeze(-2).unsqueeze(-2)
            )
        assert src_length_mask is not None, (
            "Either position_mask or src_length_mask must be provided."
        )
        return src_length_mask.unsqueeze(-2).unsqueeze(-2)

    def get_sel(
        self,
        t: torch.Tensor,
        w: torch.Tensor,
    ) -> tuple["CVMMSel", torch.Tensor]:
        """Get the selection indices and values for the mixture of experts.

        Parameters
        ----------
        t : torch.Tensor
            The input tensor for which the selection is computed. It is expected to have
            shape [..., N, C], where N is the number of elements and C is the
            projection size multiplied by the number of heads.
        w : torch.nn.Parameter
            The weight parameter used for the selection. It is expected to have shape
            [H, C], where H is the number of heads and C is the projection size.

        Returns
        -------
        tuple[CVMMSel, torch.Tensor]
            A tuple containing:
            - CVMMSel: An object containing the selection indices and values.
            - torch.Tensor: The raw selection values before applying the top-k.

        """
        sel = F.linear(t, w).float()
        sel = sel_raw = sel.view(*sel.shape[:-1], self.n_heads, -1)
        sel = sel.sigmoid()

        with torch.no_grad():
            if self.attn_heads_expert_dropout > 0 and self.training:
                mask = torch.rand_like(sel) < self.attn_heads_expert_dropout
                sel2 = sel.masked_fill(mask, float("-inf"))
            else:
                sel2 = sel
            _, sel_index = sel2.topk(
                self.attn_heads_n_active_experts,
                dim=-1,
                sorted=False,
            )
        sel_val = torch.gather(sel, -1, sel_index)

        sel_index_shifted = (
            torch.arange(self.n_heads, device=sel_index.device, dtype=sel_index.dtype)
            * self.attn_heads_n_experts
        ).unsqueeze(-1) + sel_index
        return self.cvmm_prepare_sel2(
            sel_index_shifted.flatten(-2, -1),
            sel_val,
        ), sel_raw

    def _masked_projection(  # noqa: PLR0914
        self,
        x: torch.Tensor,
        scoring_tensor: torch.Tensor,
        scoring_function: torch.nn.Parameter,
        linear_projector: torch.nn.Parameter,
        *,
        is_v: bool = True,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Perform a masked projection of the input tensor.

        This function projects the input tensor to the selection space using a linear
        projector and applies a selection mechanism based on the scoring tensor and
        scoring function. It selects the top-k elements from the selection scores and
        applies the corresponding expert projections to the input tensor.

        Parameters
        ----------
        x : torch.Tensor
            The input tensor to project. The expected shape is
            (batch_size, seq_length, d_model)
        scoring_tensor : torch.Tensor
            The tensor used to compute the selection scores. The expected shape is
            (batch_size, seq_length, d_model)
        scoring_function : torch.nn.Parameter
            The linear layer used to compute the selection scores. The expected shape is
            (attn_heads_n_experts, n_heads, d_model)
        linear_projector : torch.nn.Parameter
            The linear layer used to project the input tensor. The expected shape is
            (attn_heads_n_experts, n_heads, attn_heads_projection_size)
        is_v : bool, optional
            A flag indicating whether the `x` tensor to project refers to the values
            tensor or the output one.

        Returns
        -------
        tuple[torch.Tensor, torch.Tensor]
            A tuple containing:
            - The projected tensor whose expected shape is
                (batch_size, seq_length, d_model)
            - The selection scores whose expected shape is
                (batch_size, seq_length, n_heads, attn_heads_n_experts)

        """
        # Define the output shape
        output_shape = (*x.shape[:-1], self.n_heads * self.attn_heads_projection_size)

        # Get the appropriate experts assignments counter based on the boolean flag
        current_experts_assignments_counter = (
            self.experts_assignments_counter_v
            if is_v
            else self.experts_assignments_counter_o
        )

        # Project the input tensor to the selection space
        sel = F.linear(
            scoring_tensor,
            scoring_function.view(
                self.n_heads * self.attn_heads_n_experts,
                -1,
            ),
        ).float()
        # Reshape the selection scores to separate number of attention heads from the
        # number of experts per attention head dimension
        sel = sel_raw = sel.view(
            *sel.shape[:-1],
            self.n_heads,
            self.attn_heads_n_experts,
        )
        # Apply the sigmoid activation function to the selection scores
        # to ensure they are in the range [0, 1].
        sel = sel.sigmoid()
        with torch.no_grad():
            # Apply dropout to the selection scores if required
            if self.attn_heads_expert_dropout > 0 and self.training:
                mask = torch.rand_like(sel) < self.attn_heads_expert_dropout
                sel2 = sel.masked_fill(mask, float("-inf"))
            else:
                sel2 = sel
            # Select the top-k elements from the selection scores of each head's experts
            # (on the last dimension)
            _, sel_index = sel2.topk(
                self.attn_heads_n_active_experts,
                dim=-1,
                sorted=False,
            )
        # Handle the case where all experts are active vs. subset selection
        if self.attn_heads_n_active_experts == self.attn_heads_n_experts:
            # When all experts are active, create indices for all experts for each token
            batch_size, seq_length, n_heads = sel.shape[:3]
            # For gathering, we need indices for all experts for each head
            gather_indices = (
                torch.arange(
                    self.attn_heads_n_experts,
                    dtype=sel_index.dtype,
                    device=sel_index.device,
                )
                .unsqueeze(0)
                .unsqueeze(0)
                .unsqueeze(0)
                .expand(
                    batch_size,
                    seq_length,
                    n_heads,
                    -1,
                )
            )
            # Create binary mask that selects all experts for all tokens and heads
            sel_index_binary = torch.ones_like(gather_indices)
        else:
            # Convert the tensor `sel_index` to its one-hot version for subset selection
            sel_index_onehot = F.one_hot(
                sel_index,
                num_classes=self.attn_heads_n_experts,
            ).to(dtype=sel_index.dtype)
            # Sum along the n_active_experts dimension to get the final mask
            sel_index_binary = sel_index_onehot.sum(dim=-2)
            gather_indices = sel_index

        # Gather the selection scores for the selected indices
        sel_val = torch.gather(sel, -1, gather_indices)

        # Flatten the input tensor on all dimensions except the one of the matrix
        # multiplication with the linear projector
        x = x.flatten(0, 1)
        # Initialize the result tensor to have on the first rank the total number of
        # heads, the middle rank to have the token dimensions multiplied, and the output
        # of the linear projector as last. This is done for convenience in the loop
        # below.
        y = torch.zeros(
            (self.n_heads, *x.shape[:-1], self.attn_heads_projection_size),
            device=x.device,
            dtype=x.dtype,
        )
        # Flatten on the tokens dimensions (batch_size and seq_length) and reshape to
        # have as first rank product of the number of attention heads and the number of
        # experts per attention head
        sel_index_reshaped = sel_index_binary.view(
            self.n_heads,
            self.attn_heads_n_experts,
            -1,
        )
        # Reshape to have as first rank the number of attention heads, and second the
        # number of experts per attention head, last the tokens dimensions
        sel_val = sel_val.flatten(0, 1).view(
            self.n_heads,
            self.attn_heads_n_experts,
            -1,
        )
        with torch.no_grad():
            self.get_experts_scores_statistics(sel_val, is_v=is_v)

        # Optimized computation using broadcasting and einsum
        expert_mask = sel_index_reshaped.bool()

        # Expand tensors for efficient broadcasting
        expert_mask_expanded = expert_mask.unsqueeze(-1)
        sel_val_expanded = sel_val.unsqueeze(-1)

        # Broadcast input across heads and experts
        x_broadcast = x.unsqueeze(0).unsqueeze(0)
        x_broadcast = x_broadcast.expand(
            self.n_heads,
            self.attn_heads_n_experts,
            -1,
            -1,
        )

        # Apply expert mask
        x_masked = x_broadcast * expert_mask_expanded

        # Reshape linear projector for efficient computation
        projector_reshaped = linear_projector.transpose(0, 1)

        # Efficient batch matrix multiplication using einsum
        expert_outputs = torch.einsum(
            "hetd,hedp->hetp",
            x_masked,
            projector_reshaped,
        )

        # Apply selection scores - handle subset case properly
        if self.attn_heads_n_active_experts == self.attn_heads_n_experts:
            # All experts case
            weighted_outputs = expert_outputs * sel_val_expanded
        else:
            # Subset case: need to scatter sel_val into the right positions
            # sel_val has shape [n_heads, n_active_experts, seq_len]
            # Create tensor to hold all expert weights
            batch_size, seq_length = sel.shape[:2]
            sel_weights_full = torch.zeros(
                batch_size,
                seq_length,
                self.n_heads,
                self.attn_heads_n_experts,
                dtype=sel_val.dtype,
                device=sel_val.device,
            )
            # Use sel_index_binary as mask to place sel_val values in correct positions
            # sel_index_binary has shape [batch_size, seq_length, n_heads, n_experts]
            # sel_val needs to be reshaped to flatten properly
            sel_val_flat = sel_val.permute(2, 0, 1).flatten()
            sel_weights_full.masked_scatter_(
                sel_index_binary.bool(),
                sel_val_flat,
            )
            # Reshape back to [n_heads, n_experts, seq_len] for computation
            sel_weights_reshaped = sel_weights_full.permute(2, 3, 0, 1).contiguous()
            sel_weights_reshaped = sel_weights_reshaped.view(
                self.n_heads,
                self.attn_heads_n_experts,
                -1,
            )
            weighted_outputs = expert_outputs * sel_weights_reshaped.unsqueeze(-1)

        # Sum across experts dimension
        y = weighted_outputs.sum(dim=1)

        # Update expert assignment counters (vectorized)
        with torch.no_grad():
            assignments = expert_mask.sum(dim=-1)
            for head_idx in range(self.n_heads):
                for expert_idx in range(self.attn_heads_n_experts):
                    current_experts_assignments_counter[head_idx][expert_idx] = int(
                        assignments[head_idx, expert_idx].item(),
                    )
        # Reshape the output to respect what the function is supposed to return
        return y.view(output_shape), sel_raw

    def attend(
        self,
        v: torch.Tensor,
        k: torch.Tensor,
        q: torch.Tensor,
        mask: torch.Tensor | None,
        *,
        output_attentions: bool = False,
    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Perform attention with rotary positional encoding.

        Parameters
        ----------
        pos_offset : int
            Offset to apply to the sequence dimension, typically used for sliding window
            attention mechanisms.
        v : torch.Tensor
            The value tensor which is expected to have shape
            (batch_size, seq_len, n_heads, attn_heads_projection_size).
        k : torch.Tensor
            The key tensor which is expected to have shape
            (batch_size, seq_len, n_heads, attn_heads_projection_size).
        q : torch.Tensor
            The query tensor which is expected to have shape
            (batch_size, seq_len, n_heads, attn_heads_projection_size).
        mask : torch.Tensor | None
            Currently unused
        output_attentions : bool, optional
            Whether to return the attention probabilities along with the attention
            output. Default is False.

        Returns
        -------
        torch.Tensor | tuple[torch.Tensor, torch.Tensor, torch.Tensor]
            Tensor containing the attention output, shape
            (batch_size, seq_len, n_heads, attn_heads_projection_size), or, if
            `output_attentions` is True, it returns a tuple with three tensors: the
            attention output, the logsumexp of each row of the matrix QK^T * scaling
            (e.g., log of the softmax normalization factor), and the output of the
            softmax. The shape of the second tensor in the tuple is
            (batch_size, n_heads, seq_len). The shape of the third tensor in the tuple
            is (batch_size, n_heads, seq_len, seq_len).

        """
        if mask is not None:
            msg = "Masking attention is not implemented for SwitchHead yet."
            raise NotImplementedError(msg)
        if self.use_flash_attn:
            # NOTE: In order to use the packed flash attention implementation, we need
            # to stack them along the third dimension (dim=2), so that the resulting
            # `qkv` tensor has shape (batch_size, seq_len, 3, n_heads, projection_size).
            qkv = torch.stack([q, k, v], dim=2)
            return cast(
                "torch.Tensor | tuple[torch.Tensor, torch.Tensor, torch.Tensor]",
                self.flash_attn_qkvpacked_func(
                    qkv=qkv,
                    # NOTE: The dropout is applied to the attention scores later
                    dropout_p=0,
                    # TODO(anon): Check for this
                    softmax_scale=None,
                    causal=True,
                    # NOTE: If not (-1, -1), implements sliding window local attention.
                    window_size=(-1, -1),
                    # TODO(anon): Check for this
                    softcap=0,
                    # NOTE: (nheads,) or (batch_size, nheads), fp32. A bias of
                    # (-alibi_slope * |i - j|) is added to the attention score of query
                    # i and key j.
                    alibi_slopes=None,
                    # NOTE: Whether to use the deterministic implementation of the
                    # backward pass, which is slightly slower and uses more memory. The
                    # forward pass is always deterministic.
                    deterministic=False,
                    return_attn_probs=output_attentions,
                ),
            )
        return F.scaled_dot_product_attention(
            q,
            k,
            v,
            ~mask if mask is not None else None,
            scale=1.0,
        )

    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 no selection
            history is available, it returns 0.

        """
        loss = torch.Tensor(0)
        if self.sel_hist:
            for i in range(len(self.sel_hist[0])):
                loss += entropy_reg(
                    torch.stack([val[i] for val in self.sel_hist], dim=-3).flatten(
                        -4,
                        -3,
                    ),
                    -3,
                )
        self.sel_hist = []
        return loss

    def _get_attention_scores(  # noqa: PLR0913
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        pos_offset: int,
        mask: AttentionMask | None,
        *,
        output_attentions: bool = False,
    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Get the attention scores for the SwitchHead layer.

        Parameters
        ----------
        q : torch.Tensor
            The query tensor, expected shape is
            (batch_size, sequence_length, attn_heads_projection_size * n_heads).
        k : torch.Tensor
            The key tensor, expected shape is
            (batch_size, sequence_length, attn_heads_projection_size * n_heads).
        v : torch.Tensor
            The value tensor, expected shape is
            (batch_size, sequence_length, attn_heads_projection_size * n_heads).
        pos_offset : torch.Tensor
            The positional offset to apply to the sequence dimension, typically used for
            sliding window attention mechanisms. This is an integer value
            representing the offset to apply to the sequence length.
        mask : AttentionMask | None
            The attention mask containing position and source length masks.
        output_attentions : bool, optional
            Whether to return the attention probabilities along with the attention
            output. Default is False.

        Returns
        -------
        torch.Tensor | tuple[torch.Tensor, torch.Tensor, torch.Tensor]
            Tensor containing the attention output, shape
            (batch_size, seq_len, n_heads, attn_heads_projection_size), or, if
            `output_attentions` is True, it returns a tuple with three tensors: the
            attention output, the logsumexp of each row of the matrix QK^T * scaling
            (e.g., log of the softmax normalization factor), and the output of the
            softmax. The shape of the second tensor in the tuple is
            (batch_size, n_heads, seq_len). The shape of the third tensor in the tuple
            is (batch_size, n_heads, seq_len, seq_len).

        """
        q = q.view(
            *q.shape[:-1],
            self.n_heads,
            self.attn_heads_projection_size,
        )
        k = k.view(
            *k.shape[:-1],
            self.n_heads,
            self.attn_heads_projection_size,
        )
        v = v.view(
            *v.shape[:-1],
            self.n_heads,
            self.attn_heads_projection_size,
        )
        if self.rope_module is not None:
            q, k = self.rope_module(q, k, pos_offset or 0)
        q = self.attn_heads_dropout(q)
        return self.attend(
            v=v,
            k=k,
            q=q,
            mask=self.get_mask_tensor(v.shape[-2], mask),
            output_attentions=output_attentions,
        )

    def _forward_simplified(  # noqa: PLR0913
        self,
        q_src: torch.Tensor,
        k_src: torch.Tensor,
        v_src: torch.Tensor,
        mask: AttentionMask | None,
        kv_cache: KVCache = None,
        *,
        output_attentions: bool = False,
    ) -> tuple[torch.Tensor, KVCache, torch.Tensor | None]:
        """Forward pass of the SwitchHead layer.

        Parameters
        ----------
        q_src : torch.Tensor
            The query source tensor, expected shape is
            [batch_size, sequence_length, d_model], where `c` is the input size.
        k_src : torch.Tensor
            The key source tensor, expected shape is
            [batch_size, sequence_length, d_model], where `c` is the input size.
        v_src : torch.Tensor
            The value source tensor, expected shape is
            [batch_size, sequence_length, d_model], where `c` is the input size.
        mask : AttentionMask | None
            The attention mask containing position and source length masks.
        kv_cache : KVCache, optional
            The key-value cache for storing past key and value tensors. If provided,
            it should contain the keys "v" and "k" for the value and key tensors
            respectively. If not provided, a new cache will be created.
        output_attentions : bool, optional
            Whether to return the attention probabilities along with the attention
            output. Default is False.

        Returns
        -------
        tuple[torch.Tensor, KVCache]
            A tuple containing:
            - torch.Tensor: The output tensor after applying the attention mechanism,
              with shape [batch_size, sequence_length, d_model], where `c` is the output
              size.
            - KVCache: The updated key-value cache containing the new key and value
              tensors. If `kv_cache` was provided, it will be updated with the new
              key and value tensors. If not provided, a new cache will be created
            and returned.

        """
        assert self.attn_heads_n_active_experts == self.attn_heads_n_experts == 1, (
            "This method is only for the case where the number of experts is 1 (equal"
            " to the number of active experts)."
        )
        # Obtain the position offset in case the query and key tensors have
        # different sequence lengths.
        pos_offset = q_src.shape[1] - k_src.shape[1]
        assert pos_offset >= 0
        # Initialize the attention probabilities for returning
        attn_probs: torch.Tensor | None = None

        # Get the scale for the attention scores, i.e., the square root of the
        # projection size. This is used to scale the query and key tensors before
        # computing the attention scores.
        scale = cast("torch.Tensor", self.scale)

        # Project the query and key tensors to the attention heads projection size
        # NOTE: These now have the following shape
        # shape is (batch size, sequence length, projection_size * n_heads)
        q = cast("torch.Tensor", self.q(q_src))
        k = cast("torch.Tensor", self.k(k_src))

        # Scale the query and key tensors by the square root of the projection size
        # NOTE: These in-place modification of `q` and `k` are not allowed by PyTorch
        # Autograd when FSDP wrapped using amp_bf16
        q = q * scale.type_as(q)  # noqa: PLR6104
        k = k * scale.type_as(k)  # noqa: PLR6104
        # Get the selection scores for the output gate
        o_gate = F.sigmoid(
            F.linear(
                q_src,
                self.sel_o.view(
                    self.attn_heads_n_experts * self.n_heads,
                    self.input_size,
                ),
            ),
        )
        # Project the value tensor to the attention heads projection size
        v = F.linear(v_src, self.v)
        # Add to the key-value cache if provided.
        if kv_cache is not None:
            v = torch.cat([kv_cache["v"], v], dim=-2) if "v" in kv_cache else v
            k = torch.cat([kv_cache["k"], k], dim=-2) if "k" in kv_cache else k
            kv_cache = {"v": v, "k": k}
        # Operate the attention module using the FlashAttention and RoPe
        results = self._get_attention_scores(
            q=q,
            k=k,
            v=v,
            pos_offset=pos_offset,
            mask=mask,
            output_attentions=output_attentions,
        )
        # If the attention probabilities are returned, unpack them
        if isinstance(results, tuple):
            res, attn_probs, _softmax = results
        else:
            res = results
        # Scale the attention out by the gating matrix
        # NOTE: In-place modification of `res` is not allowed by PyTorch Autograd
        res = res * o_gate[..., None]  # noqa: PLR6104
        # Apply the output projection
        out = F.linear(res.flatten(-2), self.o)

        return out, kv_cache, attn_probs

    def _forward_cvmm_kernel(  # noqa: PLR0913, PLR0914
        self,
        q_src: torch.Tensor,
        k_src: torch.Tensor,
        v_src: torch.Tensor,
        mask: AttentionMask | None,
        kv_cache: KVCache = None,
        *,
        output_attentions: bool = False,
    ) -> tuple[torch.Tensor, KVCache, torch.Tensor | None]:
        """Forward pass of the SwitchHead layer.

        Parameters
        ----------
        q_src : torch.Tensor
            The query source tensor, expected shape is
            [batch_size, sequence_length, d_model], where `c` is the input size.
        k_src : torch.Tensor
            The key source tensor, expected shape is
            [batch_size, sequence_length, d_model], where `c` is the input size.
        v_src : torch.Tensor
            The value source tensor, expected shape is
            [batch_size, sequence_length, d_model], where `c` is the input size.
        mask : AttentionMask | None
            The attention mask containing position and source length masks.
        kv_cache : KVCache, optional
            The key-value cache for storing past key and value tensors. If provided,
            it should contain the keys "v" and "k" for the value and key tensors
            respectively. If not provided, a new cache will be created.
        output_attentions : bool, optional
            Whether to return the attention probabilities along with the attention
            output. Default is False.

        Returns
        -------
        tuple[torch.Tensor, KVCache]
            A tuple containing:
            - torch.Tensor: The output tensor after applying the attention mechanism,
              with shape [batch_size, sequence_length, d_model], where `c` is the output
              size.
            - KVCache: The updated key-value cache containing the new key and value
              tensors. If `kv_cache` was provided, it will be updated with the new
              key and value tensors. If not provided, a new cache will be created
            and returned.

        """
        # Obtain the position offset in case the query and key tensors have
        # different sequence lengths.
        pos_offset = q_src.shape[1] - k_src.shape[1]
        assert pos_offset >= 0

        assert isinstance(self.v, torch.nn.Parameter), (
            "The `v` parameter must be a torch.nn.Parameter for the custom kernel "
            "forward pass.",
        )
        assert isinstance(self.o, torch.nn.Parameter), (
            "The `o` parameter must be a torch.nn.Parameter for the custom kernel "
            "forward pass.",
        )

        # Get the scale for the attention scores, i.e., the square root of the
        # projection size. This is used to scale the query and key tensors before
        # computing the attention scores.
        scale = cast("torch.Tensor", self.scale)

        # Project the query and key tensors to the attention heads projection size
        # NOTE: These now have the following shape
        # shape is (batch size, sequence length, projection_size * n_heads)
        q = cast("torch.Tensor", self.q(q_src))
        k = cast("torch.Tensor", self.k(k_src))

        # Scale the query and key tensors by the square root of the projection size
        # NOTE: These in-place modification of `q` and `k` are not allowed by PyTorch
        # Autograd when FSDP wrapped using amp_bf16
        q = q * scale.type_as(q)  # noqa: PLR6104
        k = k * scale.type_as(k)  # noqa: PLR6104

        attn_probs: torch.Tensor | None = None
        v_sel, v_sel_r = self.get_sel(
            k_src,
            self.sel_v.view(
                self.n_heads * self.attn_heads_n_experts,
                self.input_size,
            ),
        )
        o_sel, o_sel_r = self.get_sel(
            q_src,
            self.sel_o.view(
                self.n_heads * self.attn_heads_n_experts,
                self.input_size,
            ),
        )
        if self.training and self.n_repeats > 1:
            self.sel_hist.append((o_sel_r, v_sel_r))
        # The output selection indices are calculated from the current state and are
        # also used for projecting "q". Hoerver, that projection needs to create
        # multiple copies for the different heads. Here we already have the heads,
        # but we have to create copies for the top-k elements. We can calculate that
        # from the reduction weight. We also want to compute not only the weighted
        # average between the top-k elements, but also
        # of the different heads. So reshape the reduction weight accordingly.
        output_selector_output_index = o_sel.out_index
        assert output_selector_output_index is not None, (
            "Output index should not be None",
        )
        output_selector_reduction_weight = o_sel.reduction_weight
        assert output_selector_reduction_weight is not None, (
            "Reduction weight should not be None",
        )
        o_sel.sel_index = (
            output_selector_output_index // output_selector_reduction_weight.shape[-1]
        )
        o_sel.reduction_weight = output_selector_reduction_weight.flatten(-2)
        # Projecting the values only on the active experts paths and reshaping to be
        # consistent with query and key tensors

        v = self.cvmm(
            v_src,
            v_sel,
            self.v.view(
                self.n_heads * self.attn_heads_n_experts,
                self.input_size,
                self.attn_heads_projection_size,
            ),
        )

        v = v.view(*v.shape[:-2], self.n_heads * self.attn_heads_projection_size)
        results = self._get_attention_scores(
            q=q,
            k=k,
            v=v,
            pos_offset=pos_offset,
            mask=mask,
            output_attentions=output_attentions,
        )
        if isinstance(results, tuple):
            res, attn_probs, _softmax = results
        else:
            res = results

        out = self.cvmm(
            res,
            o_sel,
            self.o.view(
                self.n_heads * self.attn_heads_n_experts,
                self.output_size,
                self.attn_heads_projection_size,
            ).transpose(-2, -1),
        )

        if kv_cache is not None:
            v = torch.cat([kv_cache["v"], v], dim=-2) if "v" in kv_cache else v
            k = torch.cat([kv_cache["k"], k], dim=-2) if "k" in kv_cache else k
            kv_cache = {"v": v, "k": k}

        return out, kv_cache, attn_probs

    def _forward_no_kernel(  # noqa: PLR0913
        self,
        q_src: torch.Tensor,
        k_src: torch.Tensor,
        v_src: torch.Tensor,
        mask: AttentionMask | None,
        kv_cache: KVCache = None,
        *,
        output_attentions: bool = False,
    ) -> tuple[torch.Tensor, KVCache, torch.Tensor | None]:
        """Forward pass of the SwitchHead layer.

        Parameters
        ----------
        q_src : torch.Tensor
            The query source tensor, expected shape is
            [batch_size, sequence_length, d_model], where `c` is the input size.
        k_src : torch.Tensor
            The key source tensor, expected shape is
            [batch_size, sequence_length, d_model], where `c` is the input size.
        v_src : torch.Tensor
            The value source tensor, expected shape is
            [batch_size, sequence_length, d_model], where `c` is the input size.
        mask : AttentionMask | None
            The attention mask containing position and source length masks.
        kv_cache : KVCache, optional
            The key-value cache for storing past key and value tensors. If provided,
            it should contain the keys "v" and "k" for the value and key tensors
            respectively. If not provided, a new cache will be created.
        output_attentions : bool, optional
            Whether to return the attention probabilities along with the attention
            output. Default is False.

        Returns
        -------
        tuple[torch.Tensor, KVCache]
            A tuple containing:
            - torch.Tensor: The output tensor after applying the attention mechanism,
              with shape [batch_size, sequence_length, d_model], where `c` is the output
              size.
            - KVCache: The updated key-value cache containing the new key and value
              tensors. If `kv_cache` was provided, it will be updated with the new
              key and value tensors. If not provided, a new cache will be created
            and returned.

        """
        # Obtain the position offset in case the query and key tensors have
        # different sequence lengths.
        pos_offset = q_src.shape[1] - k_src.shape[1]
        assert pos_offset >= 0

        assert isinstance(
            self.v,
            torch.nn.Parameter,
        ), (
            "The `v` parameter must be a torch.nn.Parameter layer for the no kernel "
            "forward pass.",
        )

        assert isinstance(
            self.o,
            torch.nn.Parameter,
        ), (
            "The `o` parameter must be a torch.nn.Parameter layer for the no kernel "
            "forward pass.",
        )

        # Get the scale for the attention scores, i.e., the square root of the
        # projection size. This is used to scale the query and key tensors before
        # computing the attention scores.
        scale = cast("torch.Tensor", self.scale)

        # Project the query and key tensors to the attention heads projection size
        # NOTE: These now have the following shape
        # shape is (batch size, sequence length, projection_size * n_heads)
        q = cast("torch.Tensor", self.q(q_src))
        k = cast("torch.Tensor", self.k(k_src))

        # Scale the query and key tensors by the square root of the projection size
        # NOTE: These in-place modification of `q` and `k` are not allowed by PyTorch
        # Autograd when FSDP wrapped using amp_bf16
        q = q * scale.type_as(q)  # noqa: PLR6104
        k = k * scale.type_as(k)  # noqa: PLR6104

        # Initialize the attention probabilities for returning
        attn_probs: torch.Tensor | None = None

        # Get the projected values only for the selected experts
        v, v_sel_r = self._masked_projection(
            x=v_src,
            scoring_tensor=k_src,
            scoring_function=self.sel_v,
            linear_projector=self.v,
        )
        results = self._get_attention_scores(
            q=q,
            k=k,
            # TODO(anon): Figure out why the dtype has changed
            v=v.to(dtype=q.dtype),
            pos_offset=pos_offset,
            mask=mask,
            output_attentions=output_attentions,
        )
        if isinstance(results, tuple):
            res, attn_probs, _softmax = results
        else:
            res = results

        # Get the projected results only for the selected experts
        out, o_sel_r = self._masked_projection(
            x=res.flatten(-2, -1),
            scoring_tensor=q_src,
            scoring_function=self.sel_o,
            linear_projector=self.o,
            is_v=False,
        )
        if self.training and self.n_repeats > 1:
            self.sel_hist.append((o_sel_r, v_sel_r))

        if kv_cache is not None:
            v = torch.cat([kv_cache["v"], v], dim=-2) if "v" in kv_cache else v
            k = torch.cat([kv_cache["k"], k], dim=-2) if "k" in kv_cache else k
            kv_cache = {"v": v, "k": k}

        return out, kv_cache, attn_probs

    def forward(
        self,
        x: torch.Tensor,
        mask: AttentionMask | None,
        kv_cache: KVCache = None,
        *,
        output_attentions: bool = False,
    ) -> tuple[torch.Tensor, KVCache, torch.Tensor | None]:
        """Forward pass of the SwitchHead layer.

        Parameters
        ----------
        x : torch.Tensor
            The maindata path tensor, expected shape is
            [batch_size, sequence_length, d_model], where `c` is the input size.
        mask : AttentionMask | None
            The attention mask containing position and source length masks.
        kv_cache : KVCache, optional
            The key-value cache for storing past key and value tensors. If provided,
            it should contain the keys "v" and "k" for the value and key tensors
            respectively. If not provided, a new cache will be created.
        output_attentions : bool, optional
            Whether to return the attention probabilities along with the attention
            output. Default is False.

        Returns
        -------
        tuple[torch.Tensor, KVCache]
            A tuple containing:
            - torch.Tensor: The output tensor after applying the attention mechanism,
              with shape [batch_size, sequence_length, d_model], where `c` is the output
              size.
            - KVCache: The updated key-value cache containing the new key and value
              tensors. If `kv_cache` was provided, it will be updated with the new
              key and value tensors. If not provided, a new cache will be created
            and returned.

        """
        # Reset the experts assignments counters
        self._reset_experts_assignments_counters()
        # Apply Layer Normalization to the input tensor
        x_norm = self.input_ln(x)
        # Operate the correct MHA module depending on the configuration

        # NOTE: If n_repeats == 1, we default to a pre-layernorm transformer
        # 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.attn_heads_n_active_experts == self.attn_heads_n_experts == 1:
            y = self._forward_simplified(
                q_src=x_norm,
                k_src=x_norm,
                v_src=x_norm if (self.use_peri_ln or self.n_repeats == 1) else x,
                mask=mask,
                kv_cache=kv_cache,
                output_attentions=output_attentions,
            )
        elif self.cvmm_kernel:
            y = self._forward_cvmm_kernel(
                q_src=x_norm,
                k_src=x_norm,
                v_src=x_norm if (self.use_peri_ln or self.n_repeats == 1) else x,
                mask=mask,
                kv_cache=kv_cache,
                output_attentions=output_attentions,
            )
        else:
            y = self._forward_no_kernel(
                q_src=x_norm,
                k_src=x_norm,
                v_src=x_norm if (self.use_peri_ln or self.n_repeats == 1) else x,
                mask=mask,
                kv_cache=kv_cache,
                output_attentions=output_attentions,
            )
        # Apply PeriLN to the output tensor
        if self.use_peri_ln:
            y = self.peri_ln_out(y)
        return y
