import torch


def get_occupancy_function(F, sharpness):
    """
    Arguments:
    ----------
        F: Tensor with size BxNx1 containing the distance function for each
        point to its closest primitive
        sharpness: Tensor with size BxMx2 containing the sharpness parameter
                   for each primitive for the positive and negative points

    Returns:
    --------
        F: Tensor with size BxNxM containing the value of the occupancy
           function of each primitive for the N points.
    """
    B, N, _ = F.shape
    _, M, _ = sharpness.shape

    mask_inside = (F < F.new_tensor(1.0)).float()
    F_bar = 1.0 - F

    s_in = sharpness[:, :, 0]
    s_out = sharpness[:, :, 1]
    F_inside = torch.sigmoid(s_in.unsqueeze(1) * F_bar) * mask_inside
    F_out = torch.sigmoid(s_out.unsqueeze(1) * F_bar) * (1-mask_inside)
    F_bar = F_inside + F_out
    assert F_bar.shape == (B, N, 1)

    return F_bar