import torch
import math


def bbox_iou(box1, box2, x1y1x2y2=True, giou=False, diou=False, ciou=False, eps=1e-9):
    """Compute the IoU of box1 to box2. box1 is 4, box2 is nx4

    :param box1: box1
    :param box2: box2
    :param x1y1x2y2: xyxy or xywh
    :param giou: giou
    :param diou: diou
    :param ciou: ciou
    :param eps: 1e-9
    :return: iou or giou, diou, ciou
    """

    box2 = box2.T

    # Get the coordinates of bounding boxes
    if x1y1x2y2:  # x1, y1, x2, y2 = box1
        box1_x1, box1_y1, box1_x2, box1_y2 = box1[0], box1[1], box1[2], box1[3]
        box2_x1, box2_y1, box2_x2, box2_y2 = box2[0], box2[1], box2[2], box2[3]
    else:  # transform from xywh to xyxy
        box1_x1, box1_x2 = box1[0] - box1[2] / 2, box1[0] + box1[2] / 2
        box1_y1, box1_y2 = box1[1] - box1[3] / 2, box1[1] + box1[3] / 2
        box2_x1, box2_x2 = box2[0] - box2[2] / 2, box2[0] + box2[2] / 2
        box2_y1, box2_y2 = box2[1] - box2[3] / 2, box2[1] + box2[3] / 2

    # Intersection area
    inter = (torch.min(box1_x2, box2_x2) - torch.max(box1_x1, box2_x1)).clamp(0) * \
            (torch.min(box1_y2, box2_y2) - torch.max(box1_y1, box2_y1)).clamp(0)

    # Union area
    box1_w, box1_h = box1_x2 - box1_x1, box1_y2 - box1_y1 + eps
    box2_w, box2_h = box2_x2 - box2_x1, box2_y2 - box2_y1 + eps
    union = box1_w * box1_h + box2_w * box2_h - inter + eps

    iou = inter / union
    if giou or diou or ciou:
        convex_width = torch.max(box1_x2, box2_x2) - torch.min(box1_x1, box2_x1)   # convex width
        convex_height = torch.max(box1_y2, box2_y2) - torch.min(box1_y1, box2_y1)  # convex height
        if ciou or diou:  # Distance or Complete IoU https://arxiv.org/abs/1911.08287v1
            convex_diagonal_squared = convex_width ** 2 + convex_height ** 2 + eps
            center_distance_squared = ((box2_x1 + box2_x2 - box1_x1 - box1_x2) ** 2 +
                                       (box2_y1 + box2_y2 - box1_y1 - box1_y2) ** 2) / 4
            if diou:
                return iou - center_distance_squared / convex_diagonal_squared
            if ciou:  # https://github.com/Zzh-tju/DIoU-SSD-pytorch/blob/master/utils/box/box_utils.py#L47
                upsilon = (4 / math.pi ** 2) * torch.pow(torch.atan(box2_w / box2_h) - torch.atan(box1_w / box1_h), 2)
                with torch.no_grad():
                    alpha = upsilon / ((1 + eps) - iou + upsilon)
                return iou - (center_distance_squared / convex_diagonal_squared + upsilon * alpha)
        else:  # giou https://arxiv.org/pdf/1902.09630.pdf
            convex_area = convex_width * convex_height + eps
            return iou - (convex_area - union) / convex_area

    return iou