import torch


def bondmatrix_masking(cutoff):
    """Returns function that returns masking matrix that is 1 where
    distance smaller cutoff, 0 when larger cutoff
    Args:
        cutoff: distance Float
    Returns:
        Callable that returns  masking matrix from a distance matrix
    """

    # @jax.profiler.trace_function
    # @jit
    def create_bondmatrix(dr):
        mask = torch.where(torch.logical_and(dr > 0.1, dr < cutoff), 1.0, 0.0)
        return mask

    return create_bondmatrix