import torch


def binary_cross_entropy_loss(score, gt_score, ratio=-1):
    """Get the binary cross entropy loss
    Args:
        score: multi-dim tensor, predicted score
        gt_score: multi-dim tensor, ground-truth score, 1 is positive, 0 is negative
        ratio: the ratio number of negative and positive samples, if -1, we use all negative samples
    """
    # positive samples
    margin = 1e-10
    mask = gt_score == 1
    score_pos = score[mask]  # 1D tensor
    score_pos = score_pos[score_pos > margin]
    #score_pos[score_pos < margin] = score_pos[score_pos < margin] * 0 + margin
    loss_pos = 0 - torch.log(score_pos)
    num_pos = loss_pos.size(0)

    # negative samples
    mask = gt_score == 0
    score_neg = score[mask]  # 1D tensor
    score_neg = 1 - score_neg
    score_neg = score_neg[score_neg > margin]
    #score_neg[score_neg < margin] = score_neg[score_neg < margin] * 0 + margin
    loss_neg = 0 - torch.log(score_neg)
    num_neg = loss_neg.size(0)
    if ratio > 0:
        loss_neg = loss_neg.sort(descending=True)[0]
        num_neg = min(loss_neg.size(0), int(ratio*num_pos))
        loss_neg = loss_neg[0:num_neg]

    if num_neg > 0 and num_pos > 0:
        weight_pos = num_neg / (num_pos + num_neg)
        weight_neg = num_pos / (num_pos + num_neg)
        loss = weight_pos * loss_pos.mean() + weight_neg * loss_neg.mean()
    elif num_neg == 0 and num_pos > 0:
        loss = loss_pos.mean()
    elif num_neg > 0 and num_pos == 0:
        loss = loss_neg.mean()
    else:
        loss = score.mean() * 0  # torch.Tensor([0]).to(score.device)

    return loss