import cv2
import torch
import numpy as np

def loc_accuracy(outputs, labels, gt_boxes, bboxes, iou_threshold=0.5):
    if outputs is not None:
        _, pred = torch.topk(outputs, k=1, dim=1, largest=True, sorted=True)
        pred = pred.t()
        correct = pred.eq(labels.view(1, -1).expand_as(pred))
        wrongs = [c == 0 for c in correct.cpu().numpy()][0]

    batch_size = len(gt_boxes)
    gt_known, top1 = 0., 0.
    for i, (gt_box, bbox) in enumerate(zip(gt_boxes, bboxes)):
        iou_score = iou(gt_box, bbox)

        if iou_score >= iou_threshold:
            gt_known += 1.
            if outputs is not None and not wrongs[i]:
                top1 += 1.

    gt_loc = gt_known / batch_size
    top1_loc = top1 / batch_size
    return gt_loc, top1_loc

def iou(box1, box2):
    """box: (xmin, ymin, xmax, ymax)"""
    box1_xmin, box1_ymin, box1_xmax, box1_ymax = box1
    box2_xmin, box2_ymin, box2_xmax, box2_ymax = box2

    inter_xmin = max(box1_xmin, box2_xmin)
    inter_ymin = max(box1_ymin, box2_ymin)
    inter_xmax = min(box1_xmax, box2_xmax)
    inter_ymax = min(box1_ymax, box2_ymax)

    inter_area = (inter_xmax - inter_xmin + 1) * (inter_ymax - inter_ymin + 1)
    box1_area = (box1_xmax - box1_xmin + 1) * (box1_ymax - box1_ymin + 1)
    box2_area = (box2_xmax - box2_xmin + 1) * (box2_ymax - box2_ymin + 1)

    iou = inter_area / (box1_area + box2_area - inter_area).float()
    return iou.item()