import torch


def calculate_tp_troch(pred_boxes, pred_scores, gt_boxes, gt_difficult=None, iou_thresh=0.5):
    """
        calculate tp/fp for all predicted bboxes for one class of one image.
        对于匹配到同一gt的不同bboxes，让score最高tp = 1，其它的tp = 0
    Args:
        pred_boxes: Tensor[N, 4], 某张图片中某类别的全部预测框的坐标 (x0, y0, x1, y1)
        pred_scores: Tensor[N, 1], 某张图片中某类别的全部预测框的score
        gt_boxes: Tensor[M, 4], 某张图片中某类别的全部gt的坐标 (x0, y0, x1, y1)
        gt_difficult: Tensor[M, 1], 某张图片中某类别的gt中是否为difficult目标的值
        iou_thresh: iou 阈值

    Returns:
        gt_num: 某张图片中某类别的gt数量
        tp_list: 记录某张图片中某类别的预测框是否为tp的情况
        confidence_score: 记录某张图片中某类别的预测框的score值 (与tp_list相对应)
    """
    if gt_boxes.numel() == 0:
        return 0, [], []

    # 若无对应的boxes，则 tp 为空
    if pred_boxes.numel() == 0:
        return len(gt_boxes), [], []

    if gt_difficult is None:
        gt_difficult = torch.zeros((len(gt_boxes), 1)).to(gt_boxes)

    # 否则计算所有预测框与gt之间的iou
    ious = pred_boxes.new_zeros((len(gt_boxes), len(pred_boxes)))
    for i in range(len(gt_boxes)):
        gb = gt_boxes[i]
        area_pb = (pred_boxes[:, 2] - pred_boxes[:, 0]) * (pred_boxes[:, 3] - pred_boxes[:, 1])
        area_gb = (gb[2] - gb[0]) * (gb[3] - gb[1])
        xx1 = pred_boxes[:, 0].clamp(min=gb[0].item())  # [N-1,]
        yy1 = pred_boxes[:, 1].clamp(min=gb[1].item())
        xx2 = pred_boxes[:, 2].clamp(max=gb[2].item())
        yy2 = pred_boxes[:, 3].clamp(max=gb[3].item())
        inter = (xx2 - xx1).clamp(min=0) * (yy2 - yy1).clamp(min=0)  # [N-1,]
        ious[i] = inter / (area_pb + area_gb - inter)
    # 每个预测框的最大iou所对应的gt记为其匹配的gt
    max_ious, max_ious_idx = ious.max(dim=0)

    not_difficult_gt_mask = gt_difficult == 0
    gt_num = not_difficult_gt_mask.sum().item()
    if gt_num == 0:
        return 0, [], []

    # 保留 max_iou 中属于 非difficult 目标的预测框，即应该去掉与 difficult gt 相匹配的预测框，不参与p-r计算
    # 如果去掉与 difficult gt 对应的iou分数后，候选框的最大iou依然没有发生改变，则可认为此候选框不与difficult gt相匹配，应该保留
    not_difficult_pb_mask = (ious[not_difficult_gt_mask].max(dim=0)[0] == max_ious)
    max_ious, max_ious_idx = max_ious[not_difficult_pb_mask], max_ious_idx[not_difficult_pb_mask]
    if max_ious_idx.numel() == 0:
        return gt_num, [], []

    confidence_score = pred_scores.view(-1)[not_difficult_pb_mask]
    tp_list = torch.zeros_like(max_ious)
    for i in max_ious_idx[max_ious > iou_thresh].unique():
        gt_mask = (max_ious > iou_thresh) * (max_ious_idx == i)
        idx = (confidence_score * gt_mask.float()).argmax()
        tp_list[idx] = 1

    return gt_num, tp_list.tolist(), confidence_score.tolist()