import torch

from mmdet3d.ops.iou3d import boxes_iou3d_gpu
from . import box_torch_ops


def bbox_overlaps_2d(bboxes1, bboxes2, mode='iou', is_aligned=False):
    """Calculate overlap between two set of bboxes.

    If ``is_aligned`` is ``False``, then calculate the ious between each bbox
    of bboxes1 and bboxes2, otherwise the ious between each aligned pair of
    bboxes1 and bboxes2.

    Args:
        bboxes1 (Tensor): shape (m, 4) in <x1, y1, x2, y2> format.
        bboxes2 (Tensor): shape (n, 4) in <x1, y1, x2, y2> format.
            If is_aligned is ``True``, then m and n must be equal.
        mode (str): "iou" (intersection over union) or iof (intersection over
            foreground).

    Returns:
        ious(Tensor): shape (m, n) if is_aligned == False else shape (m, 1)

    Example:
        >>> bboxes1 = torch.FloatTensor([
        >>>     [0, 0, 10, 10],
        >>>     [10, 10, 20, 20],
        >>>     [32, 32, 38, 42],
        >>> ])
        >>> bboxes2 = torch.FloatTensor([
        >>>     [0, 0, 10, 20],
        >>>     [0, 10, 10, 19],
        >>>     [10, 10, 20, 20],
        >>> ])
        >>> bbox_overlaps(bboxes1, bboxes2)
        tensor([[0.5238, 0.0500, 0.0041],
                [0.0323, 0.0452, 1.0000],
                [0.0000, 0.0000, 0.0000]])

    Example:
        >>> empty = torch.FloatTensor([])
        >>> nonempty = torch.FloatTensor([
        >>>     [0, 0, 10, 9],
        >>> ])
        >>> assert tuple(bbox_overlaps(empty, nonempty).shape) == (0, 1)
        >>> assert tuple(bbox_overlaps(nonempty, empty).shape) == (1, 0)
        >>> assert tuple(bbox_overlaps(empty, empty).shape) == (0, 0)
    """

    assert mode in ['iou', 'iof']

    rows = bboxes1.size(0)
    cols = bboxes2.size(0)
    if is_aligned:
        assert rows == cols

    if rows * cols == 0:
        return bboxes1.new(rows, 1) if is_aligned else bboxes1.new(rows, cols)

    if is_aligned:
        lt = torch.max(bboxes1[:, :2], bboxes2[:, :2])  # [rows, 2]
        rb = torch.min(bboxes1[:, 2:], bboxes2[:, 2:])  # [rows, 2]

        wh = (rb - lt).clamp(min=0)  # [rows, 2]
        overlap = wh[:, 0] * wh[:, 1]
        area1 = (bboxes1[:, 2] - bboxes1[:, 0]) * (
            bboxes1[:, 3] - bboxes1[:, 1])

        if mode == 'iou':
            area2 = (bboxes2[:, 2] - bboxes2[:, 0]) * (
                bboxes2[:, 3] - bboxes2[:, 1])
            ious = overlap / (area1 + area2 - overlap)
        else:
            ious = overlap / area1
    else:
        lt = torch.max(bboxes1[:, None, :2], bboxes2[:, :2])  # [rows, cols, 2]
        rb = torch.min(bboxes1[:, None, 2:], bboxes2[:, 2:])  # [rows, cols, 2]

        wh = (rb - lt).clamp(min=0)  # [rows, cols, 2]
        overlap = wh[:, :, 0] * wh[:, :, 1]
        area1 = (bboxes1[:, 2] - bboxes1[:, 0]) * (
            bboxes1[:, 3] - bboxes1[:, 1])

        if mode == 'iou':
            area2 = (bboxes2[:, 2] - bboxes2[:, 0]) * (
                bboxes2[:, 3] - bboxes2[:, 1])
            ious = overlap / (area1[:, None] + area2 - overlap)
        else:
            ious = overlap / (area1[:, None])

    return ious


def bbox_overlaps_3d(bboxes1, bboxes2, mode='iou'):
    '''

    :param bboxes1: Tensor, shape (N, 7) [x, y, z, h, w, l, ry]
    :param bboxes2: Tensor, shape (M, 7) [x, y, z, h, w, l, ry]
    :param mode: mode (str): "iou" (intersection over union) or
            iof (intersection over foreground).
    :return: iou: (M, N) not support aligned mode currently
    '''
    # TODO: check the input dimension meanings,
    #  this is inconsistent with that in bbox_overlaps_nearest_3d
    return boxes_iou3d_gpu(bboxes1, bboxes2, mode)


def bbox_overlaps_nearest_3d(bboxes1, bboxes2, mode='iou', is_aligned=False):
    '''
    :param bboxes1: Tensor, shape (N, 7) [x, y, z, h, w, l, ry]?
    :param bboxes2: Tensor, shape (M, 7) [x, y, z, h, w, l, ry]?
    :param mode: mode (str): "iou" (intersection over union) or iof
            (intersection over foreground).
    :return: iou: (M, N) not support aligned mode currently
    rbboxes: [N, 5(x, y, xdim, ydim, rad)] rotated bboxes
    '''
    # TODO: check the input dimension meanings,
    # this is inconsistent with that in bbox_overlaps_3d
    rbboxes1_np = bboxes1.index_select(
        dim=-1, index=bboxes1.new_tensor([0, 1, 3, 4, 6]).long())
    rbboxes2_np = bboxes2.index_select(
        dim=-1, index=bboxes1.new_tensor([0, 1, 3, 4, 6]).long())

    # Change the bboxes to bev
    # box conversion and iou calculation in torch version on CUDA
    # is 10x faster than that in numpy version
    bboxes1_bv = box_torch_ops.rbbox2d_to_near_bbox(rbboxes1_np)
    bboxes2_bv = box_torch_ops.rbbox2d_to_near_bbox(rbboxes2_np)
    ret = bbox_overlaps_2d(
        bboxes1_bv, bboxes2_bv, mode=mode, is_aligned=is_aligned)
    return ret