import torch


def nms(bboxes, scores, threshold=0.5):
    """
        Performs non-maximum suppression (NMS) on the boxes according
        to their intersection-over-union (IoU).

        NMS iteratively removes lower scoring boxes which have an
        IoU greater than iou_threshold with another (higher scoring)
        box.

        Arguments:
            bboxes (Tensor[N, 4]): boxes to perform NMS on
            scores (Tensor[N]): scores for each one of the boxes
            threshold (float): discards all overlapping boxes with IoU < iou_threshold

        Returns:
            keep (Tensor): int64 tensor with the indices of the elements that have been kept
                by NMS, sorted in decreasing order of scores
    """
    x1 = bboxes[:, 0]
    y1 = bboxes[:, 1]
    x2 = bboxes[:, 2]
    y2 = bboxes[:, 3]
    areas = (x2 - x1) * (y2 - y1)  # [N,] 每个bbox的面积
    _, order = scores.sort(0, descending=True)  # 降序排列

    keep = []
    while order.numel() > 0:  # torch.numel()返回张量元素个数
        if order.numel() == 1:  # 保留框只剩一个
            i = order.item()
            keep.append(i)
            break
        else:
            i = order[0].item()  # 保留scores最大的那个框box[i]
            keep.append(i)

        # 计算box[i]与其余各框的IOU(思路很好)
        xx1 = x1[order[1:]].clamp(min=x1[i].item())  # [N-1,]
        yy1 = y1[order[1:]].clamp(min=y1[i].item())
        xx2 = x2[order[1:]].clamp(max=x2[i].item())
        yy2 = y2[order[1:]].clamp(max=y2[i].item())
        inter = (xx2 - xx1).clamp(min=0) * (yy2 - yy1).clamp(min=0)  # [N-1,]

        iou = inter / (areas[i] + areas[order[1:]] - inter)  # [N-1,]
        idx = (iou <= threshold).nonzero().squeeze()  # 注意此时idx为[N-1,] 而order为[N,]
        if idx.numel() == 0:
            break
        order = order[idx + 1]  # 修补索引之间的差值

    return torch.tensor(keep).long()