"""
Code adapted from https://www.pyimagesearch.com/2016/11/07/intersection-over-union-iou-for-object-detection/
"""
import numpy as np
import torch


def intersection_over_union(bb_a, bb_b):
    """Compute the intersection over union between two sets of bounding boxes

    For computing the IoU, the edges of the bounding boxes are assumed to be
    part of the object as well.

    Args:
        bb_a (BoundingBoxes): First set of bounding boxes.
        bb_b (BoundingBoxes): Second set of bounding boxes.

    Returns:
        torch.tensor: A matrix of IoU values of shape (bb_a.shape[0],
            bb_b.shape[0]). The matrix contains all pairwise IoU values
            for all bounding boxes in bb_a and bb_b.
    """
    with torch.no_grad():
        orig_encoding_a = bb_a.encoding
        orig_encoding_b = bb_b.encoding
        orig_encoding_params_a = bb_a.encoding_params
        orig_encoding_params_b = bb_b.encoding_params
        bb_a.decode()
        bb_b.decode()
        ullr_a = bb_a.coords.numpy()
        ullr_b = bb_b.coords.numpy()

        ul_x_a = np.tile(ullr_a[:, 0].reshape(-1, 1), (1, ullr_b.shape[0]))
        ul_y_a = np.tile(ullr_a[:, 1].reshape(-1, 1), (1, ullr_b.shape[0]))
        lr_x_a = np.tile(ullr_a[:, 2].reshape(-1, 1), (1, ullr_b.shape[0]))
        lr_y_a = np.tile(ullr_a[:, 3].reshape(-1, 1), (1, ullr_b.shape[0]))
        ul_x_b = np.tile(ullr_b[:, 0], (ullr_a.shape[0], 1))
        ul_y_b = np.tile(ullr_b[:, 1], (ullr_a.shape[0], 1))
        lr_x_b = np.tile(ullr_b[:, 2], (ullr_a.shape[0], 1))
        lr_y_b = np.tile(ullr_b[:, 3], (ullr_a.shape[0], 1))

        # determine the (x, y)-coordinates of the intersection rectangle
        xA = np.maximum(ul_x_a, ul_x_b)
        yA = np.maximum(ul_y_a, ul_y_b)
        xB = np.minimum(lr_x_a, lr_x_b)
        yB = np.minimum(lr_y_a, lr_y_b)

        # compute the area of intersection rectangle
        inter_areas = np.maximum(0, xB - xA + 1) * np.maximum(0, yB - yA + 1)

        # compute the areas of both sets of bounding boxes
        areas_a = (lr_x_a - ul_x_a + 1) * (lr_y_a - ul_y_a + 1)
        areas_b = (lr_x_b - ul_x_b + 1) * (lr_y_b - ul_y_b + 1)

        # compute the intersection over union
        ious = inter_areas / (areas_a + areas_b - inter_areas)

        # encode the bounding boxes back to their original encoding
        bb_a.encode(
            encoding=orig_encoding_a,
            encoding_params=orig_encoding_params_a,
        )
        bb_b.encode(
            encoding=orig_encoding_b,
            encoding_params=orig_encoding_params_b,
        )

        # return the intersection over union values
    return torch.from_numpy(ious)
