import torch
from typing import Union
from typing import Callable
import warnings


def do_activation(input_data: torch.Tensor, activation: Union[str, Callable] = "softmax") -> torch.Tensor:
    """
    This function is used to do activation for inputs.

    Args:
        input_data: the input that to be activated, in the shape [B] or [BN] or [BNHW] or [BNHWD].
        activation: can be ``"sigmoid"`` or ``"softmax"``, or a callable function. Defaults to ``"softmax"``.
            An example for callable function: ``activation = lambda x: torch.log_softmax(x)``.

    Raises:
        NotImplementedError: When input an activation name that is not implemented.
    """
    input_ndim = input_data.ndimension()
    if activation == "softmax":
        if input_ndim == 1:
            warnings.warn("input_data has only one channel, softmax ignored.")
        else:
            input_data = input_data.float().softmax(dim=1)
    elif activation == "sigmoid":
        input_data = input_data.float().sigmoid()
    elif callable(activation):
        input_data = activation(input_data)
    else:
        raise NotImplementedError("activation can only be sigmoid, softmax or a callable function.")
    return input_data