"""MuPSharedEmbedding class for muP models.

This class extends SharedEmbedding to include scaling and optional normalization.
"""

from typing import Any

from llmfoundry.models.layers import build_norm
from llmfoundry.models.layers.custom_embedding import SharedEmbedding
from torch import Tensor


class MuPSharedEmbedding(SharedEmbedding):
    """SharedEmbedding that scales its forward output for muP.

    And optionally applies normalization.
    """

    def __init__(  # noqa: PLR0913
        self,
        vocab_size: int,
        d_model: int,
        *args: Any,  # noqa: ANN401
        scale: float = 1.0,
        use_embedding_norm: bool = False,
        norm_type: str = "low_precision_layernorm",
        norm_eps: float = 1e-05,
        device: str | None = None,
        **kwargs: Any,  # noqa: ANN401
    ) -> None:
        """Initialize the MuPSharedEmbedding.

        Parameters
        ----------
        vocab_size: int
            Size of the vocabulary.
        d_model: int
            Dimension of the model.
        *args: Any
            Additional arguments for SharedEmbedding.
        scale: float
            Scaling factor for the embedding output.
        use_embedding_norm: bool
            Whether to apply normalization to the embedding output.
        norm_type: str
            Type of normalization to apply (e.g., 'layernorm').
        norm_eps: float
            Epsilon value for normalization.
        device: str | None
            Device to initialize the embedding on (e.g., 'cpu', 'cuda').
        **kwargs: Any
            Additional keyword arguments for SharedEmbedding.

        """
        super().__init__(vocab_size, d_model, *args, device=device, **kwargs)
        # Begin muP code: store scaling factor for embedding outputs
        self.scale = scale
        # End muP code
        self.norm = (
            None
            if not use_embedding_norm
            else build_norm(
                name=norm_type.lower(),
                normalized_shape=d_model,
                eps=norm_eps,
                device=device,
            )
        )

    def forward(self, input: Tensor, unembed: bool = False) -> Tensor:  # noqa: A002, FBT001, FBT002
        """Forward pass of the MuPSharedEmbedding.

        Parameters
        ----------
        input: Tensor
            Input tensor containing token indices.
        unembed: bool
            If True, the output will not be scaled and will be used for unembedding.

        Returns
        -------
        Tensor: The output tensor after applying the embedding
        and optional normalization.

        """
        out = super().forward(input, unembed)
        if not unembed:
            out = self.norm(out) if self.norm is not None else out
            # Begin muP code: scale embedding output
            out = out * self.scale  # noqa: PLR6104
            # End muP code
        return out
