import torch


def step_function_perturbation(perturb, epsilon_0, alpha=1e-4, norm_type='inf', smooth_approx=False, temperature=0.01):
    """
    Step function applied to the perturbation norm. By default, it computes the exact step function which is not
    differentiable. If a smooth approximation based on the sigmoid function is needed, set `smooth_approx=True` and set
    the `temperature` to a suitably small value.

    :param perturb: Torch Tensor with the perturbation. Can be a tensor of shape `(b, d1, ...)`, where `b` is the batch
                    size and the rest are dimensions. Can also be a single vector of shape `[d]`.
    :param epsilon_0: Radius of the smaller perturbation ball - a small positive value.
    :param alpha: Small negative offset for the step function. The step function's value is `-alpha` when the
                  perturbation norm is less than `epsilon_0`.
    :param norm_type: Type of norm: 'inf' for infinity norm, '1', '2' etc for the other types of norm.
    :param smooth_approx: Set to True to get a sigmoid-based approximation of the step function.
    :param temperature: small non-negative value that controls  the steepness of the sigmoid approximation.

    :returns: tensor of function values computed for each element in the batch. Has shape `[b]`.
    """
    assert isinstance(perturb, (torch.Tensor, torch.autograd.Variable)), ("Input 'perturb' should be of type "
                                                                          "torch.Tensor or torch.autograd.Variable")
    s = perturb.shape
    dim = 1
    if len(s) > 2:
        perturb = perturb.view(s[0], -1)    # flatten into a vector
    elif len(s) == 1:
        # single vector
        dim = None

    if norm_type == 'inf':
        norm_type = float('inf')
    elif not isinstance(norm_type, (int, float)):
        # example: norm_type = '2'
        norm_type = int(norm_type)

    norm_val = torch.linalg.vector_norm(perturb, ord=norm_type, dim=dim)
    if not smooth_approx:
        return torch.where(norm_val <= epsilon_0, -1. * alpha, 1.)
    else:
        return torch.sigmoid((1. / temperature) * (norm_val - epsilon_0)) - alpha