import math


def hard_negative_mining(loss, labels, neg_pos_ratio=3):
    """
    用于训练过程中正负例比例的限制.默认在训练时,负例数量是正例数量的三倍
    Args:
        loss (N, num_priors): the loss for each example.
        labels (N, num_priors): the labels.
        neg_pos_ratio:  正负例比例: 负例数量/正例数量
    """
    pos_mask = labels > 0
    num_pos = pos_mask.long().sum(dim=1, keepdim=True)
    num_neg = num_pos * neg_pos_ratio

    loss[pos_mask] = -math.inf  # 无穷
    # 两次sort 找出元素在排序中的位置
    _, indexes = loss.sort(dim=1, descending=True)  # descending 降序 ,返回 value,index
    _, orders = indexes.sort(dim=1)
    neg_mask = orders < num_neg  # loss 降序排, 背景为-无穷, 选择排前num_neg的 负无穷,也就是 背景
    return pos_mask | neg_mask