import torch


def get_weights(labels, threshold):
    """
    Take a mini-batch of label vectors (i.e., signal envelopes) and a
    threshold, and return a vector that is 0 if 0 < label <= threshold, and
    1 otherwise, thus allowing us to make the signal parts below the
    threshold fuzzy for the computation of our metrics.

    Args:
        labels: A mini-batch of label vectors (signal envelopes / amplitudes)
        threshold: A number to use as a threshold on the signal amplitude

    Returns:
        A weights tensor (0 for fuzzy, 1 for non-fuzzy) that can be
        multiplied with the real and predicted labels before evaluating the
        loss and Hamming distance
    """

    weights = torch.eq(torch.gt(labels, threshold) *
                       torch.lt(labels, 1.2 * threshold), 0)

    return weights.float()