
import numpy as np


def calculate_iou(box1, box2):

    # Calculate the (x1, y1, x2, y2) coordinates of the intersection of box1 and box2. Calculate its Area.
    # box1 = [-3, 8.5, 3, 11.5]
    # box2 = [-3, 9.5, 3, 12.5]
    # box1 = [1086.84, 156.24, 1181.62, 319.12]
    # box2 = [1078.333357, 159.086347, 1193.771014, 322.239107]

    xi1 = max(box1[0], box2[0])
    yi1 = max(box1[1], box2[1])
    xi2 = min(box1[2], box2[2])
    yi2 = min(box1[3], box2[3])
    inter_area = max((xi2 - xi1), 0) * max((yi2 - yi1), 0)  # Max keeps into account not overlapping box

    # Calculate the Union area by using Formula: Union(A,B) = A + B - Inter(A,B)
    box1_area = (box1[2] - box1[0]) * (box1[3] - box1[1])
    box2_area = (box2[2] - box2[0]) * (box2[3] - box2[1])
    union_area = box1_area + box2_area - inter_area

    # compute the IoU
    iou = inter_area / union_area

    return iou


def get_iou_matrix(boxes, boxes_gt):
    """
    Get IoU matrix between predicted and ground truth boxes
    Dim: (boxes, boxes_gt)
    """
    iou_matrix = np.zeros((len(boxes), len(boxes_gt)))
    for idx, box in enumerate(boxes):
        for idx_gt, box_gt in enumerate(boxes_gt):
            iou_matrix[idx, idx_gt] = calculate_iou(box, box_gt)
    return iou_matrix


def get_iou_matches(boxes, boxes_gt, iou_min=0.3):
    """From 2 sets of boxes and a minimum threshold, compute the matching indices for IoU matches"""

    matches = []
    used = []
    if not boxes or not boxes_gt:
        return []
    confs = [box[4] for box in boxes]

    indices = list(np.argsort(confs))
    for idx in indices[::-1]:
        box = boxes[idx]
        ious = []
        for idx_gt, box_gt in enumerate(boxes_gt):
            iou = calculate_iou(box, box_gt)
            ious.append(iou)
        idx_gt_max = int(np.argmax(ious))
        if (ious[idx_gt_max] >= iou_min) and (idx_gt_max not in used):
            matches.append((int(idx), idx_gt_max))
            used.append(idx_gt_max)
    return matches


def get_iou_matches_matrix(boxes, boxes_gt, thresh):
    """From 2 sets of boxes and a minimum threshold, compute the matching indices for IoU matchings"""

    iou_matrix = get_iou_matrix(boxes, boxes_gt)
    if not iou_matrix.size:
        return []

    matches = []
    iou_max = np.max(iou_matrix)
    while iou_max > thresh:
        # Extract the indeces of the max
        args_max = np.unravel_index(np.argmax(iou_matrix, axis=None), iou_matrix.shape)
        matches.append(args_max)
        iou_matrix[args_max[0], :] = 0
        iou_matrix[:, args_max[1]] = 0
        iou_max = np.max(iou_matrix)
    return matches


def reorder_matches(matches, boxes, mode='left_rigth'):
    """
    Reorder a list of (idx, idx_gt) matches based on position of the detections in the image
    ordered_boxes = (5, 6, 7, 0, 1, 4, 2, 4)
    matches = [(0, x), (2,x), (4,x), (3,x), (5,x)]
    Output --> [(5, x), (0, x), (3, x), (2, x), (5, x)]
    """

    assert mode == 'left_right'

    # Order the boxes based on the left-right position in the image and
    ordered_boxes = np.argsort([box[0] for box in boxes])  # indices of boxes ordered from left to right
    matches_left = [int(idx) for (idx, _) in matches]

    return [matches[matches_left.index(idx_boxes)] for idx_boxes in ordered_boxes if idx_boxes in matches_left]
