#!/usr/bin/env python3

# 3D IoU caculate code for 3D object detection
# Kent 2018/12

import numpy as np
from scipy.spatial import ConvexHull
from numpy import *


class _3DIoU():
    """ Calculate the Interception over the Union between a pair of 3D boxes.
    """

    def __init__(self, box1, box2):
        """ Each box consists of a list containing:
        box --> [box_size: tuple of (length (X axis), wide (Y axis), height (Z axis)),
                 heading_angle: rad scalar, clockwise from pos x axis,
                 center: tuple of (x, y, z)]
        """
        self.box1 = box1
        self.box2 = box2

        self.corners1 = self.get_3d_box(
            self.box1[0], self.box1[1], self.box1[2])
        self.corners2 = self.get_3d_box(
            self.box2[0], self.box2[1], self.box2[2])

        (self.IOU_3d, self.IOU_2d) = self.box3d_iou(
            self.corners2, self.corners1)

    def polygon_clip(self, subjectPolygon, clipPolygon):
        """ Clip a polygon with another polygon.
        Ref: https://rosettacode.org/wiki/Sutherland-Hodgman_polygon_clipping#Python
        Args:
          subjectPolygon: a list of (x,y) 2d points, any polygon.
          clipPolygon: a list of (x,y) 2d points, has to be *convex*
        Note:
          **points have to be counter-clockwise ordered**
        Return:
          a list of (x,y) vertex point for the intersection polygon.
        """
        def inside(p):
            return(cp2[0]-cp1[0])*(p[1]-cp1[1]) > (cp2[1]-cp1[1])*(p[0]-cp1[0])

        def computeIntersection():
            dc = [cp1[0] - cp2[0], cp1[1] - cp2[1]]
            dp = [s[0] - e[0], s[1] - e[1]]
            n1 = cp1[0] * cp2[1] - cp1[1] * cp2[0]
            n2 = s[0] * e[1] - s[1] * e[0]
            n3 = 1.0 / (dc[0] * dp[1] - dc[1] * dp[0])
            return [(n1*dp[0] - n2*dc[0]) * n3, (n1*dp[1] - n2*dc[1]) * n3]

        outputList = subjectPolygon
        cp1 = clipPolygon[-1]

        for clipVertex in clipPolygon:
            cp2 = clipVertex
            inputList = outputList
            outputList = []
            s = inputList[-1]

            for subjectVertex in inputList:
                e = subjectVertex
                if inside(e):
                    if not inside(s):
                        outputList.append(computeIntersection())
                    outputList.append(e)
                elif inside(s):
                    outputList.append(computeIntersection())
                s = e
            cp1 = cp2
            if len(outputList) == 0:
                return None
        return(outputList)

    def poly_area(self, x, y):
        """ Ref: http://stackoverflow.com/questions/24467972/calculate-area-of-polygon-given-x-y-coordinates """
        return 0.5*np.abs(np.dot(x, np.roll(y, 1))-np.dot(y, np.roll(x, 1)))

    def convex_hull_intersection(self, p1, p2):
        """ Compute area of two convex hull's intersection area.
            p1,p2 are a list of (x,y) tuples of hull vertices.
            return a list of (x,y) for the intersection and its volume
        """
        inter_p = self.polygon_clip(p1, p2)
        if inter_p is not None:
            hull_inter = ConvexHull(inter_p, qhull_options='QJ')
            return inter_p, hull_inter.volume
        else:
            return None, 0.0

    def box3d_vol(self, corners):
        ''' corners: (8,3) no assumption on axis direction '''
        a = np.sqrt(np.sum((corners[0, :] - corners[1, :])**2))
        b = np.sqrt(np.sum((corners[1, :] - corners[2, :])**2))
        c = np.sqrt(np.sum((corners[0, :] - corners[4, :])**2))
        return a*b*c

    def is_clockwise(self, p):
        x = p[:, 0]
        y = p[:, 1]
        return np.dot(x, np.roll(y, 1))-np.dot(y, np.roll(x, 1)) > 0

    def box3d_iou(self, corners1, corners2):
        ''' Compute 3D bounding box IoU.
        Input:
            corners1: numpy array (8,3), assume up direction is negative Y
            corners2: numpy array (8,3), assume up direction is negative Y
        Output:
            iou: 3D bounding box IoU
            iou_2d: bird's eye view 2D bounding box IoU
        todo (kent): add more description on corner points' orders.
        '''
        # corner points are in counter clockwise order
        rect1 = [(corners1[i, 0], corners1[i, 2]) for i in range(3, -1, -1)]
        rect2 = [(corners2[i, 0], corners2[i, 2]) for i in range(3, -1, -1)]

        area1 = self.poly_area(np.array(rect1)[:, 0], np.array(rect1)[:, 1])
        area2 = self.poly_area(np.array(rect2)[:, 0], np.array(rect2)[:, 1])

        inter, inter_area = self.convex_hull_intersection(rect1, rect2)
        iou_2d = inter_area/(area1+area2-inter_area)
        ymax = min(corners1[0, 1], corners2[0, 1])
        ymin = max(corners1[4, 1], corners2[4, 1])

        inter_vol = inter_area * max(0.0, ymax-ymin)

        vol1 = self.box3d_vol(corners1)
        vol2 = self.box3d_vol(corners2)
        iou = inter_vol / (vol1 + vol2 - inter_vol)
        return iou, iou_2d

    # ----------------------------------
    # Helper functions for evaluation
    # ----------------------------------

    def get_3d_box(self, box_size, heading_angle, center):
        ''' Calculate 3D bounding box corners from its parameterization.
        Input:
            box_size: tuple of (length,wide,height)
            heading_angle: rad scalar, clockwise from pos x axis
            center: tuple of (x,y,z)
        Output:
            corners_3d: numpy array of shape (8,3) for 3D box cornders
        '''
        def roty(t):
            c = np.cos(t)
            s = np.sin(t)
            return np.array([[c,  0,  s],
                            [0,  1,  0],
                            [-s, 0,  c]])

        R = roty(heading_angle)
        l, h, w = box_size
        x_corners = [l/2, l/2, -l/2, -l/2, l/2, l/2, -l/2, -l/2]
        y_corners = [h/2, h/2, h/2, h/2, -h/2, -h/2, -h/2, -h/2]
        z_corners = [w/2, -w/2, -w/2, w/2, w/2, -w/2, -w/2, w/2]
        corners_3d = np.dot(R, np.vstack([x_corners, y_corners, z_corners]))
        corners_3d[0, :] = corners_3d[0, :] + center[0]
        corners_3d[1, :] = corners_3d[1, :] + center[1]
        corners_3d[2, :] = corners_3d[2, :] + center[2]
        corners_3d = np.transpose(corners_3d)
        return corners_3d
