"""Activation functions for Transducer."""

import torch
from packaging.version import parse as V


def get_activation(
    activation_type: str,
    ftswish_threshold: float = -0.2,
    ftswish_mean_shift: float = 0.0,
    hardtanh_min_val: int = -1.0,
    hardtanh_max_val: int = 1.0,
    leakyrelu_neg_slope: float = 0.01,
    smish_alpha: float = 1.0,
    smish_beta: float = 1.0,
    softplus_beta: float = 1.0,
    softplus_threshold: int = 20,
    swish_beta: float = 1.0,
) -> torch.nn.Module:
    """Return activation function.

    Args:
        activation_type: Activation function type.
        ftswish_threshold: Threshold value for FTSwish activation formulation.
        ftswish_mean_shift: Mean shifting value for FTSwish activation formulation.
        hardtanh_min_val: Minimum value of the linear region range for HardTanh.
        hardtanh_max_val: Maximum value of the linear region range for HardTanh.
        leakyrelu_neg_slope: Negative slope value for LeakyReLU activation formulation.
        smish_alpha: Alpha value for Smish activation fomulation.
        smish_beta: Beta value for Smish activation formulation.
        softplus_beta: Beta value for softplus activation formulation in Mish.
        softplus_threshold: Values above this revert to a linear function in Mish.
        swish_beta: Beta value for Swish variant formulation.

    Returns:
        : Activation function.

    """
    torch_version = V(torch.__version__)

    activations = {
        "ftswish": (
            FTSwish,
            {"threshold": ftswish_threshold, "mean_shift": ftswish_mean_shift},
        ),
        "hardtanh": (
            torch.nn.Hardtanh,
            {"min_val": hardtanh_min_val, "max_val": hardtanh_max_val},
        ),
        "leaky_relu": (torch.nn.LeakyReLU, {"negative_slope": leakyrelu_neg_slope}),
        "mish": (
            Mish,
            {
                "softplus_beta": softplus_beta,
                "softplus_threshold": softplus_threshold,
                "use_builtin": torch_version >= V("1.9"),
            },
        ),
        "relu": (torch.nn.ReLU, {}),
        "selu": (torch.nn.SELU, {}),
        "smish": (Smish, {"alpha": smish_alpha, "beta": smish_beta}),
        "swish": (
            Swish,
            {"beta": swish_beta, "use_builtin": torch_version >= V("1.8")},
        ),
        "tanh": (torch.nn.Tanh, {}),
        "identity": (torch.nn.Identity, {}),
    }

    act_func, act_args = activations[activation_type]

    return act_func(**act_args)


class FTSwish(torch.nn.Module):
    """Flatten-T Swish activation definition.

    FTSwish(x) = x * sigmoid(x) + threshold
                  where FTSwish(x) < 0 = threshold

    Reference: https://arxiv.org/abs/1812.06247

    Args:
        threshold: Threshold value for FTSwish activation formulation. (threshold < 0)
        mean_shift: Mean shifting value for FTSwish activation formulation.
                       (applied only if != 0, disabled by default)

    """

    def __init__(self, threshold: float = -0.2, mean_shift: float = 0) -> None:
        super().__init__()

        assert threshold < 0, "FTSwish threshold parameter should be < 0."

        self.threshold = threshold
        self.mean_shift = mean_shift

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Forward computation."""
        x = (x * torch.sigmoid(x)) + self.threshold
        x = torch.where(x >= 0, x, torch.tensor([self.threshold], device=x.device))

        if self.mean_shift != 0:
            x.sub_(self.mean_shift)

        return x


class Mish(torch.nn.Module):
    """Mish activation definition.

    Mish(x) = x * tanh(softplus(x))

    Reference: https://arxiv.org/abs/1908.08681.

    Args:
        softplus_beta: Beta value for softplus activation formulation.
                         (Usually 0 > softplus_beta >= 2)
        softplus_threshold: Values above this revert to a linear function.
                         (Usually 10 > softplus_threshold >= 20)
        use_builtin: Whether to use PyTorch activation function if available.

    """

    def __init__(
        self,
        softplus_beta: float = 1.0,
        softplus_threshold: int = 20,
        use_builtin: bool = False,
    ) -> None:
        super().__init__()

        if use_builtin:
            self.mish = torch.nn.Mish()
        else:
            self.tanh = torch.nn.Tanh()
            self.softplus = torch.nn.Softplus(
                beta=softplus_beta, threshold=softplus_threshold
            )

            self.mish = lambda x: x * self.tanh(self.softplus(x))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Forward computation."""
        return self.mish(x)


class Smish(torch.nn.Module):
    """Smish activation definition.

    Smish(x) = (alpha * x) * tanh(log(1 + sigmoid(beta * x)))
                 where alpha > 0 and beta > 0

    Reference: https://www.mdpi.com/2079-9292/11/4/540/htm.

    Args:
        alpha: Alpha value for Smish activation fomulation.
                 (Usually, alpha = 1. If alpha <= 0, set value to 1).
        beta: Beta value for Smish activation formulation.
                (Usually, beta = 1. If beta <= 0, set value to 1).

    """

    def __init__(self, alpha: float = 1.0, beta: float = 1.0) -> None:
        super().__init__()

        self.tanh = torch.nn.Tanh()

        self.alpha = alpha if alpha > 0 else 1
        self.beta = beta if beta > 0 else 1

        self.smish = lambda x: (self.alpha * x) * self.tanh(
            torch.log(1 + torch.sigmoid((self.beta * x)))
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Forward computation."""
        return self.smish(x)


class Swish(torch.nn.Module):
    """Swish activation definition.

    Swish(x) = (beta * x) * sigmoid(x)
                 where beta = 1 defines standard Swish activation.

    References:
        https://arxiv.org/abs/2108.12943 / https://arxiv.org/abs/1710.05941v1.
        E-swish variant: https://arxiv.org/abs/1801.07145.

    Args:
        beta: Beta parameter for E-Swish.
                (beta >= 1. If beta < 1, use standard Swish).
        use_builtin: Whether to use PyTorch function if available.

    """

    def __init__(self, beta: float = 1.0, use_builtin: bool = False) -> None:
        super().__init__()

        self.beta = beta

        if beta > 1:
            self.swish = lambda x: (self.beta * x) * torch.sigmoid(x)
        else:
            if use_builtin:
                self.swish = torch.nn.SiLU()
            else:
                self.swish = lambda x: x * torch.sigmoid(x)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Forward computation."""
        return self.swish(x)