#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
import torch
from maskrcnn_benchmark.layers import nms as _box_nms

from .bounding_box import BoxList
from maskrcnn_benchmark.structures.segmentation_mask import SegmentationMask
import numpy as np
import shapely
from shapely.geometry import Polygon,MultiPoint

def boxlist_nms(boxlist, nms_thresh, max_proposals=-1, score_field="score"):
    """
    Performs non-maximum suppression on a boxlist, with scores specified
    in a boxlist field via score_field.

    Arguments:
        boxlist(BoxList)
        nms_thresh (float)
        max_proposals (int): if > 0, then only the top max_proposals are kept
            after non-maxium suppression
        score_field (str)
    """
    if nms_thresh <= 0:
        return boxlist
    mode = boxlist.mode
    boxlist = boxlist.convert("xyxy")
    boxes = boxlist.bbox
    score = boxlist.get_field(score_field)
    keep = _box_nms(boxes, score, nms_thresh)
    if max_proposals > 0:
        keep = keep[:max_proposals]
    boxlist = boxlist[keep]
    return boxlist.convert(mode)


def remove_small_boxes(boxlist, min_size):
    """
    Only keep boxes with both sides >= min_size

    Arguments:
        boxlist (Boxlist)
        min_size (int)
    """
    # TODO maybe add an API for querying the ws / hs
    xywh_boxes = boxlist.convert("xywh").bbox
    _, _, ws, hs = xywh_boxes.unbind(dim=1)
    keep = ((ws >= min_size) & (hs >= min_size)).nonzero().squeeze(1)
    return boxlist[keep]


# implementation from https://github.com/kuangliu/torchcv/blob/master/torchcv/utils/box.py
# with slight modifications
def boxlist_iou(boxlist1, boxlist2):
    """Compute the intersection over union of two set of boxes.
    The box order must be (xmin, ymin, xmax, ymax).

    Arguments:
      box1: (BoxList) bounding boxes, sized [N,4].
      box2: (BoxList) bounding boxes, sized [M,4].

    Returns:
      (tensor) iou, sized [N,M].

    Reference:
      https://github.com/chainer/chainercv/blob/master/chainercv/utils/bbox/bbox_iou.py
    """
    if boxlist1.size != boxlist2.size:
        raise RuntimeError(
            "boxlists should have same image size, got {}, {}".format(
                boxlist1, boxlist2
            )
        )

    # N = len(boxlist1)
    # M = len(boxlist2)

    area1 = boxlist1.area()
    area2 = boxlist2.area()

    box1, box2 = boxlist1.bbox, boxlist2.bbox

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

    TO_REMOVE = 1

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

    iou = inter / (area1[:, None] + area2 - inter)
    return iou

# def boxlist_polygon_iou(target, proposal):
#     """Compute the intersection over union of two set of boxes.
#     The box order must be (xmin, ymin, xmax, ymax).

#     Arguments:
#       box1: (BoxList) bounding boxes, sized [N,4].
#       box2: (BoxList) bounding boxes, sized [M,4].

#     Returns:
#       (tensor) iou, sized [N,M].

#     Reference:
#       https://github.com/chainer/chainercv/blob/master/chainercv/utils/bbox/bbox_iou.py
#     """
#     if target.size != proposal.size:
#         raise RuntimeError(
#             "boxlists should have same image size, got {}, {}".format(
#                 target, proposal
#             )
#         )
#     target_polygon = target.get_field("masks").to_np_polygon()
#     proposal_polygon = proposal.get_field("masks").to_np_polygon()
#     print(target_polygon)
#     print(proposal_polygon)
#     polygon_points1 = target_polygon[0].reshape(-1, 2)
#     poly1 = Polygon(polygon_points1).convex_hull
#     polygon_points2 = proposal_polygon[0].reshape(-1, 2)
#     poly2 = Polygon(polygon_points2).convex_hull
#     union_poly = np.concatenate((polygon_points1, polygon_points2))
#     if not poly1.intersects(poly2): # this test is fast and can accelerate calculation
#         iou = 0
#     else:
#         try:
#             inter_area = poly1.intersection(poly2).area
#             #union_area = poly1.area + poly2.area - inter_area
#             union_area = MultiPoint(union_poly).convex_hull.area
#             if union_area == 0:
#                 return 0
#             iou = float(inter_area) / union_area
#         except shapely.geos.TopologicalError:
#             print('shapely.geos.TopologicalError occured, iou set to 0')
#             iou = 0
#     return iou


# TODO redundant, remove
def _cat(tensors, dim=0):
    """
    Efficient version of torch.cat
    avoids a copy if there is only a single element in a list
    """
    assert isinstance(tensors, (list, tuple))
    if len(tensors) == 1:
        return tensors[0]
    return torch.cat(tensors, dim)

def _cat_mask(masks):
    polygons_cat = []
    size = masks[0].size
    for mask in masks:
        polygons = mask.get_polygons()
        polygons_cat.extend(polygons)
    masks_cat = SegmentationMask(polygons_cat, size)
    return masks_cat


def cat_boxlist(bboxes):
    """
    Concatenates a list of BoxList (having the same image size) into a
    single BoxList

    Arguments:
        bboxes (list[BoxList])
    """
    # if bboxes is None:
    #     return None
    # if bboxes[0] is None:
    #     bboxes = [bboxes[1]
    assert isinstance(bboxes, (list, tuple))
    assert all(isinstance(bbox, BoxList) for bbox in bboxes)

    size = bboxes[0].size
    assert all(bbox.size == size for bbox in bboxes)

    mode = bboxes[0].mode
    assert all(bbox.mode == mode for bbox in bboxes)

    fields = set(bboxes[0].fields())
    assert all(set(bbox.fields()) == fields for bbox in bboxes)

    cat_boxes = BoxList(_cat([bbox.bbox for bbox in bboxes], dim=0), size, mode)

    for field in fields:
        if field == 'masks':
            data = _cat_mask([bbox.get_field(field) for bbox in bboxes])
        else:
            data = _cat([bbox.get_field(field) for bbox in bboxes], dim=0)
        cat_boxes.add_field(field, data)

    return cat_boxes


def cat_boxlist_gt(bboxes):
    """
    Concatenates a list of BoxList (having the same image size) into a
    single BoxList

    Arguments:
        bboxes (list[BoxList])
    """
    assert isinstance(bboxes, (list, tuple))
    assert all(isinstance(bbox, BoxList) for bbox in bboxes)

    size = bboxes[0].size
    # bboxes[1].set_size(size)
    assert all(bbox.size == size for bbox in bboxes)

    mode = bboxes[0].mode
    assert all(bbox.mode == mode for bbox in bboxes)

    fields = set(bboxes[0].fields())
    assert all(set(bbox.fields()) == fields for bbox in bboxes)
    if bboxes[0].bbox.sum().item() == 0:
        cat_boxes = BoxList(bboxes[1].bbox, size, mode)
    else:
        cat_boxes = BoxList(_cat([bbox.bbox for bbox in bboxes], dim=0), size, mode)

    for field in fields:
        if bboxes[0].bbox.sum().item() == 0:
            if field == 'masks':
                data = _cat_mask([bbox.get_field(field) for bbox in bboxes[1:]])
            else:
                data = _cat([bbox.get_field(field) for bbox in bboxes[1:]], dim=0)
        else:
            if field == 'masks':
                data = _cat_mask([bbox.get_field(field) for bbox in bboxes])
            else:
                data = _cat([bbox.get_field(field) for bbox in bboxes], dim=0)
        cat_boxes.add_field(field, data)

    return cat_boxes