# modified from torchvision to also return the union
def box_iou(boxes1, boxes2):
    area1 = box_area(boxes1)
    area2 = box_area(boxes2)

    lt = torch.max(boxes1[:, None, :2], boxes2[:, :2])  # [N,M,2]
    rb = torch.min(boxes1[:, None, 2:], boxes2[:, 2:])  # [N,M,2]

    wh = (rb - lt).clamp(min=0)  # [N,M,2]
    inter = wh[:, :, 0] * wh[:, :, 1]  # [N,M]

    union = area1[:, None] + area2 - inter

    iou = inter / union
    return iou, union


def distance_iou(boxes1, boxes2):
    iou, union = box_iou(boxes1, boxes2)

    lt = torch.min(boxes1[:, None, :2], boxes2[:, :2])
    rb = torch.max(boxes1[:, None, 2:], boxes2[:, 2:])

    wh = (rb - lt).clamp(min=0)  # [N,M,2]
    c2 = wh[..., 0] ** 2 + wh[..., 1] ** 2

    b1 = (boxes1[:, 2:] + boxes1[:, :2]) / 2
    b2 = (boxes2[:, 2:] + boxes2[:, :2]) / 2
    # rho2 = torch.cdist(b1, b2, p=2) ** 2
    rho2 = ((b1[:, None, :] - b2[None, :, :]) ** 2).sum(-1)

    return 1 - iou + rho2 / c2

def distance_box_iou(boxes1, boxes2):
    return 1 - distance_iou(boxes1, boxes2)

box1 = torch.tensor([[0, 0, 100, 100], [0, 0, 50, 50], [200, 200, 300, 300], [0, 0, 25, 25]])
box2 = torch.tensor([[0, 0, 100, 100], [0, 0, 50, 50], [200, 200, 300, 300]])
print(distance_box_iou(box1, box2))

from functools import partial

import torch
from torchvision import ops


def loss_to_iou(a, b, fn):
    # Convert Loss to IoU
    return 1.0 - fn(a, b)


def gen_box(size):
    # Build proper bboxes
    xy1 = torch.rand((size, 2))
    xy2 = xy1 + torch.rand((size, 2))
    return torch.cat([xy1, xy2], axis=-1)


def cartesian_product_slow(boxes1, boxes2, fn):
    # Extemely slow method for testing cartesian product and use of the loss fn
    N = boxes1.size(0)
    M = boxes2.size(0)

    result = torch.zeros((N, M))
    for i in range(N):
        for j in range(M):
            result[i, j] = fn(boxes1[i].unsqueeze(0), boxes2[j].unsqueeze(0))

    return result


def cartesian_product_fast(boxes1, boxes2, fn):
    # Faster method for testing cartesian product and use of the loss fn
    N = boxes1.size(0)
    M = boxes2.size(0)

    bx_l = boxes1.repeat_interleave(M, dim=0)
    bx_r = boxes2.repeat((N, 1))

    result = fn(bx_l, bx_r)

    return result.view(N, M)


N = 5
M = 7

boxes1 = gen_box(N)
boxes2 = gen_box(M)

# Check that the slow and fast are equivalent
a = cartesian_product_fast(boxes1, boxes2, partial(loss_to_iou, fn=ops.distance_box_iou_loss))
b = cartesian_product_slow(boxes1, boxes2, partial(loss_to_iou, fn=ops.distance_box_iou_loss))
assert torch.allclose(a, b)

# Check the build-in Generalized IoU method against the cartesian of the equivalent loss
a = cartesian_product_fast(boxes1, boxes2, partial(loss_to_iou, fn=ops.generalized_box_iou_loss))
b = ops.generalized_box_iou(boxes1, boxes2)
assert torch.allclose(a, b)

# Check the build-in Distance IoU method against the cartesian of the equivalent loss
a = cartesian_product_fast(boxes1, boxes2, partial(loss_to_iou, fn=ops.distance_box_iou_loss))
b = ops.distance_box_iou(boxes1, boxes2)
assert torch.allclose(a, b)

# Check the build-in Complete IoU method against the cartesian of the equivalent loss
a = cartesian_product_fast(boxes1, boxes2, partial(loss_to_iou, fn=ops.complete_box_iou_loss))
b = ops.complete_box_iou(boxes1, boxes2)
assert torch.allclose(a, b)

print("OK")