r"""Feature scaling layers."""

import typing

import torch
from torch import Tensor, nn


def layer_scale(
    input: Tensor, gamma: Tensor, training: bool, *, inplace: bool | None = None
) -> Tensor:
    """Functional version of :class:`LayerScale`."""
    if inplace is True or (
        inplace is None
        and (
            not torch.jit.is_scripting()
            and not torch.jit.is_tracing()
            and not torch.compiler.is_compiling()
            and not training
            and input.is_contiguous()
        )
    ):
        return input.mul_(gamma)
    return input * gamma


class LayerScale(nn.Module):
    """LayerScale methodology used in ConvNextV2."""

    def __init__(
        self,
        dim: int,
        gamma: float | Tensor = 1e-5,
        inplace: bool | None = False,
    ) -> None:
        r"""
        Parameters
        ----------
        dim : int
            The number of features in the input tensor.
        gamma : float or Tensor
            The initial value of the scaling factor.
        inplace : bool, optional
            Whether to perform the operation in-place or not. If `None`, the operation
            is performed in-place if the input tensor is contiguous and the
            module is *not* in training mode.
        """
        super().__init__()

        self.inplace = inplace
        self.gamma = nn.Parameter(gamma * torch.ones(dim))

        self.register_load_state_dict_pre_hook(self.__class__._patch_missing_gamma)

    def _patch_missing_gamma(
        self,
        state_dict: dict[str, Tensor],
        prefix: str,
        _local_metadata: dict[str, typing.Any],
        _strict: bool,
        _missing_keys: list[str],
        _unexpected_keys: list[str],
        _error_msgs: list[str],
    ) -> None:
        key_gamma = f"{prefix}gamma"
        if key_gamma not in state_dict:
            state_dict[key_gamma] = torch.ones_like(self.gamma)

    @typing.override
    def forward(self, input: Tensor) -> Tensor:
        return layer_scale(input, self.gamma, self.training, inplace=self.inplace)


class LayerScale2d(LayerScale):
    """Variant of :class:`LayerScale` for 2D inputs with shape ``(B, C, H, W)``."""

    @typing.override
    def forward(self, input: Tensor) -> Tensor:
        return layer_scale(
            input,
            self.gamma.view(1, -1, 1, 1),
            self.training,
            inplace=self.inplace,
        )
