import numpy as np
from scipy.spatial import Delaunay
import scipy
import kitti.object3d as object3d
import torch

def get_objects_from_label(label_file):
    with open(label_file, 'r') as f:
        lines = f.readlines()
    objects = [object3d.Object3d(line) for line in lines]
    return objects


def dist_to_plane(plane, points):
    """
    Calculates the signed distance from a 3D plane to each point in a list of points
    :param plane: (a, b, c, d)
    :param points: (N, 3)
    :return: (N), signed distance of each point to the plane
    """
    a, b, c, d = plane

    points = np.array(points)
    x = points[:, 0]
    y = points[:, 1]
    z = points[:, 2]

    return (a*x + b*y + c*z + d) / np.sqrt(a**2 + b**2 + c**2)


def rotate_pc_along_y(pc, rot_angle):
    """
    params pc: (N, 3+C), (N, 3) is in the rectified camera coordinate
    params rot_angle: rad scalar
    Output pc: updated pc with XYZ rotated
    """
    cosval = np.cos(rot_angle)
    sinval = np.sin(rot_angle)
    rotmat = np.array([[cosval, -sinval], [sinval, cosval]])
    pc[:, [0, 2]] = np.dot(pc[:, [0, 2]], np.transpose(rotmat))
    return pc


def rotate_pc_along_y_torch(pc, rot_angle):
    """
    :param pc: (N, 512, 3 + C)
    :param rot_angle: (N)
    :return:
    TODO: merge with rotate_pc_along_y_torch in bbox_transform.py
    """
    cosa = torch.cos(rot_angle).view(-1, 1)  # (N, 1)
    sina = torch.sin(rot_angle).view(-1, 1)  # (N, 1)

    raw_1 = torch.cat([cosa, -sina], dim=1)  # (N, 2)
    raw_2 = torch.cat([sina, cosa], dim=1)  # (N, 2)
    R = torch.cat((raw_1.unsqueeze(dim=1), raw_2.unsqueeze(dim=1)), dim=1)  # (N, 2, 2)

    pc_temp = pc[:, :, [0, 2]]  # (N, 512, 2)

    pc[:, :, [0, 2]] = torch.matmul(pc_temp, R.permute(0, 2, 1))  # (N, 512, 2)

    return pc


def boxes3d_to_corners3d(boxes3d, rotate=True):
    """
    :param boxes3d: (N, 7) [x, y, z, h, w, l, ry]
    :param rotate:
    :return: corners3d: (N, 8, 3)
    """
    boxes_num = boxes3d.shape[0]
    h, w, l = boxes3d[:, 3], boxes3d[:, 4], boxes3d[:, 5]
    x_corners = np.array([l / 2., l / 2., -l / 2., -l / 2., l / 2., l / 2., -l / 2., -l / 2.], dtype=np.float32).T  # (N, 8)
    z_corners = np.array([w / 2., -w / 2., -w / 2., w / 2., w / 2., -w / 2., -w / 2., w / 2.], dtype=np.float32).T  # (N, 8)

    y_corners = np.zeros((boxes_num, 8), dtype=np.float32)
    y_corners[:, 4:8] = -h.reshape(boxes_num, 1).repeat(4, axis=1)  # (N, 8)

    if rotate:
        ry = boxes3d[:, 6]
        zeros, ones = np.zeros(ry.size, dtype=np.float32), np.ones(ry.size, dtype=np.float32)
        rot_list = np.array([[np.cos(ry), zeros, -np.sin(ry)],
                             [zeros,       ones,       zeros],
                             [np.sin(ry), zeros,  np.cos(ry)]])  # (3, 3, N)
        R_list = np.transpose(rot_list, (2, 0, 1))  # (N, 3, 3)

        temp_corners = np.concatenate((x_corners.reshape(-1, 8, 1), y_corners.reshape(-1, 8, 1),
                                       z_corners.reshape(-1, 8, 1)), axis=2)  # (N, 8, 3)
        rotated_corners = np.matmul(temp_corners, R_list)  # (N, 8, 3)
        x_corners, y_corners, z_corners = rotated_corners[:, :, 0], rotated_corners[:, :, 1], rotated_corners[:, :, 2]

    x_loc, y_loc, z_loc = boxes3d[:, 0], boxes3d[:, 1], boxes3d[:, 2]

    x = x_loc.reshape(-1, 1) + x_corners.reshape(-1, 8)
    y = y_loc.reshape(-1, 1) + y_corners.reshape(-1, 8)
    z = z_loc.reshape(-1, 1) + z_corners.reshape(-1, 8)

    corners = np.concatenate((x.reshape(-1, 8, 1), y.reshape(-1, 8, 1), z.reshape(-1, 8, 1)), axis=2)

    return corners.astype(np.float32)


def boxes3d_to_corners3d_torch(boxes3d, flip=False):
    """
    :param boxes3d: (N, 7) [x, y, z, h, w, l, ry]
    :return: corners_rotated: (N, 8, 3)
    """
    boxes_num = boxes3d.shape[0]
    h, w, l, ry = boxes3d[:, 3:4], boxes3d[:, 4:5], boxes3d[:, 5:6], boxes3d[:, 6:7]
    if flip:
        ry = ry + np.pi
    centers = boxes3d[:, 0:3]
    zeros = torch.cuda.FloatTensor(boxes_num, 1).fill_(0)
    ones = torch.cuda.FloatTensor(boxes_num, 1).fill_(1)

    x_corners = torch.cat([l / 2., l / 2., -l / 2., -l / 2., l / 2., l / 2., -l / 2., -l / 2.], dim=1)  # (N, 8)
    y_corners = torch.cat([zeros, zeros, zeros, zeros, -h, -h, -h, -h], dim=1)  # (N, 8)
    z_corners = torch.cat([w / 2., -w / 2., -w / 2., w / 2., w / 2., -w / 2., -w / 2., w / 2.], dim=1)  # (N, 8)
    corners = torch.cat((x_corners.unsqueeze(dim=1), y_corners.unsqueeze(dim=1), z_corners.unsqueeze(dim=1)), dim=1) # (N, 3, 8)

    cosa, sina = torch.cos(ry), torch.sin(ry)
    raw_1 = torch.cat([cosa, zeros, sina], dim=1)
    raw_2 = torch.cat([zeros, ones, zeros], dim=1)
    raw_3 = torch.cat([-sina, zeros, cosa], dim=1)
    R = torch.cat((raw_1.unsqueeze(dim=1), raw_2.unsqueeze(dim=1), raw_3.unsqueeze(dim=1)), dim=1)  # (N, 3, 3)

    corners_rotated = torch.matmul(R, corners)  # (N, 3, 8)
    corners_rotated = corners_rotated + centers.unsqueeze(dim=2).expand(-1, -1, 8)
    corners_rotated = corners_rotated.permute(0, 2, 1)
    return corners_rotated


def boxes3d_to_bev_torch(boxes3d):
    """
    :param boxes3d: (N, 7) [x, y, z, h, w, l, ry]
    :return:
        boxes_bev: (N, 5) [x1, y1, x2, y2, ry]
    """
    boxes_bev = boxes3d.new(torch.Size((boxes3d.shape[0], 5)))

    cu, cv = boxes3d[:, 0], boxes3d[:, 2]
    half_l, half_w = boxes3d[:, 5] / 2, boxes3d[:, 4] / 2
    boxes_bev[:, 0], boxes_bev[:, 1] = cu - half_l, cv - half_w
    boxes_bev[:, 2], boxes_bev[:, 3] = cu + half_l, cv + half_w
    boxes_bev[:, 4] = boxes3d[:, 6]
    return boxes_bev


def enlarge_box3d(boxes3d, extra_width):
    """
    :param boxes3d: (N, 7) [x, y, z, h, w, l, ry]
    """
    if isinstance(boxes3d, np.ndarray):
        large_boxes3d = boxes3d.copy()
    else:
        large_boxes3d = boxes3d.clone()
    large_boxes3d[:, 3:6] += extra_width * 2
    large_boxes3d[:, 1] += extra_width
    return large_boxes3d


def in_hull(p, hull):
    """
    :param p: (N, K) test points
    :param hull: (M, K) M corners of a box
    :return (N) bool
    """
    try:
        if not isinstance(hull, Delaunay):
            hull = Delaunay(hull)
        flag = hull.find_simplex(p) >= 0
    except scipy.spatial.qhull.QhullError:
        print('Warning: not a hull %s' % str(hull))
        flag = np.zeros(p.shape[0], dtype=np.bool)

    return flag


def objs_to_boxes3d(obj_list):
    boxes3d = np.zeros((obj_list.__len__(), 7), dtype=np.float32)
    for k, obj in enumerate(obj_list):
        boxes3d[k, 0:3], boxes3d[k, 3], boxes3d[k, 4], boxes3d[k, 5], boxes3d[k, 6] \
            = obj.pos, obj.h, obj.w, obj.l, obj.ry
    return boxes3d


def objs_to_scores(obj_list):
    scores = np.zeros((obj_list.__len__()), dtype=np.float32)
    for k, obj in enumerate(obj_list):
        scores[k] = obj.score
    return scores


def get_iou3d(corners3d, query_corners3d, need_bev=False):
    """	
    :param corners3d: (N, 8, 3) in rect coords	
    :param query_corners3d: (M, 8, 3)	
    :return:	
    """
    from shapely.geometry import Polygon
    A, B = corners3d, query_corners3d
    N, M = A.shape[0], B.shape[0]
    iou3d = np.zeros((N, M), dtype=np.float32)
    iou_bev = np.zeros((N, M), dtype=np.float32)

    # for height overlap, since y face down, use the negative y
    min_h_a = -A[:, 0:4, 1].sum(axis=1) / 4.0
    max_h_a = -A[:, 4:8, 1].sum(axis=1) / 4.0
    min_h_b = -B[:, 0:4, 1].sum(axis=1) / 4.0
    max_h_b = -B[:, 4:8, 1].sum(axis=1) / 4.0

    for i in range(N):
        for j in range(M):
            max_of_min = np.max([min_h_a[i], min_h_b[j]])
            min_of_max = np.min([max_h_a[i], max_h_b[j]])	
            h_overlap = np.max([0, min_of_max - max_of_min])
            if h_overlap == 0:
                continue

            bottom_a, bottom_b = Polygon(A[i, 0:4, [0, 2]].T), Polygon(B[j, 0:4, [0, 2]].T)
            if bottom_a.is_valid and bottom_b.is_valid:
                # check is valid,  A valid Polygon may not possess any overlapping exterior or interior rings.
                bottom_overlap = bottom_a.intersection(bottom_b).area
            else:
                bottom_overlap = 0.
            overlap3d = bottom_overlap * h_overlap
            union3d = bottom_a.area * (max_h_a[i] - min_h_a[i]) + bottom_b.area * (max_h_b[j] - min_h_b[j]) - overlap3d
            iou3d[i][j] = overlap3d / union3d
            iou_bev[i][j] = bottom_overlap / (bottom_a.area + bottom_b.area - bottom_overlap)

    if need_bev:
        return iou3d, iou_bev

    return iou3d


def my_compute_box_3d(center, size, heading_angle):
    R = rotz(-1*heading_angle)
    l,w,h = size
    x_corners = [-l,l,l,-l,-l,l,l,-l]
    y_corners = [w,w,-w,-w,w,w,-w,-w]
    z_corners = [h,h,h,h,-h,-h,-h,-h]
    corners_3d = np.dot(R, np.vstack([x_corners, y_corners, z_corners]))
    corners_3d[0,:] += center[0]
    corners_3d[1,:] += center[1]
    corners_3d[2,:] += center[2]
    return np.transpose(corners_3d)


def in_hull(p, hull):
    from scipy.spatial import Delaunay
    if not isinstance(hull,Delaunay):
        hull = Delaunay(hull)
    return hull.find_simplex(p)>=0


def extract_pc_in_box3d(pc, box3d):
    ''' pc: (N,3), box3d: (8,3) '''
    box3d_roi_inds = in_hull(pc[:,0:3], box3d)
    return pc[box3d_roi_inds,:], box3d_roi_inds


def data_viz(objects, pts, name='demo', dump_dir='./data_viz_dump'):
    ''' Examine and visualize KITTI data. '''
    import os
    import utils.pc_util as pc_util

    # Dump OBJ files for the point cloud
    pc = flip_axis_to_viz(pts)
    write_ply_pc(pc, os.path.join(dump_dir, 'pc_%s.obj' % name))

    # Dump OBJ files for 3D bounding boxes
    # l,w,h correspond to dx,dy,dz
    # heading angle is from +X rotating towards -Y
    # (+X is degree, -Y is 90 degrees)
    oriented_boxes = []

    for obj in objects:
        obb = flip_box_to_viz(obj)
        oriented_boxes.append(obb)
    if len(oriented_boxes) > 0:
        oriented_boxes = np.vstack(tuple(oriented_boxes))
        pc_util.write_oriented_bbox(oriented_boxes,
                                    os.path.join(dump_dir, 'obbs_%s.ply' % name))
    else:
        print('No boxes to viz')


def point_viz(pts, name='demo', dump_dir='./data_viz_dump'):
    import os

    # Dump OBJ files for the point cloud
    pc = flip_axis_to_viz(pts)
    write_ply_pc(pc, os.path.join(dump_dir, 'pc_%s.obj' % name))


def box_viz(objects, name='demo', dump_dir='./data_viz_dump'):
    import os
    import utils.pc_util as pc_util
    oriented_boxes = []

    for obj in objects:
        obb = flip_box_to_viz(obj)
        oriented_boxes.append(obb)
    if len(oriented_boxes) > 0:
        oriented_boxes = np.vstack(tuple(oriented_boxes))
        pc_util.write_oriented_bbox(oriented_boxes,
                                    os.path.join(dump_dir, 'obbs_%s.ply' % name))
    else:
        print('No boxes to viz')


def flip_box_to_viz(obj):
    '''
    :param obj: x-right, y-down, z-forward, h, w, l, ry
    :return: obj2: x-right, y-forward, z-up, l, w, h, rz
    '''
    # obj2 = np.copy(obj)
    x, y, z = obj[0], obj[2], -1 * obj[1]
    h, w, l = obj[3], obj[4], obj[5]
    rz = -1 * obj[6]
    z = z + h / 2
    obj2 = np.array([x, y, z, l, w, h, rz])
    return obj2


def single_box3d_to_corner3d(box):
    corners = boxes3d_to_corners3d(box.reshape((1, -1)))
    return corners[0]


def corer3d_to_box3d(corners):
    '''
    :param corners: numpy array [8, 3]

    coord:             box:
       ______z             6_______7
      /\                   /\     /\
    x/ \                  / .    / \
       \y                / 2\___/__\3
                        /   /  /   /
                      5/______/4  /
                       \  /   \  /
                       \ .    \ /
                       \/_____\/
                       1       0

    :return: [x, y, z, h, w, l, ry] in camera coord, y is on the bottom surface of the box
    '''
    import math
    center_x = (corners[0, 0] + corners[2, 0]) / 2
    center_z = (corners[0, 2] + corners[2, 2]) / 2
    y = corners[0, 1]
    corners_tmp = corners.copy()

    corners_tmp[:, 0] = corners_tmp[:, 0] - center_x
    corners_tmp[:, 2] = corners_tmp[:, 2] - center_z
    angle_point = (corners_tmp[0] + corners_tmp[1]) / 2
    angle = math.atan(angle_point[2] / angle_point[0])
    if angle_point[0] < 0:
        angle = angle + math.pi
    if angle > math.pi:
        angle = angle - 2 * math.pi
    elif angle < -1 * math.pi:
        angle = angle + 2 * math.pi
    angle = -1 * angle

    l = math.sqrt(math.pow(corners[1, 0] - corners[2, 0], 2) + math.pow(corners[1, 2] - corners[2, 2], 2))
    w = math.sqrt(math.pow(corners[0, 0] - corners[1, 0], 2) + math.pow(corners[0, 2] - corners[1, 2], 2))
    h = corners[0, 1] - corners[4, 1]
    box3d = np.array([center_x, y, center_z, h, w, l, angle])
    return box3d


def flip_axis_to_viz(pc):
    ''' Flip X-right,Y-down,Z-forward to X-right,Y-forward,Z-up
    Input and output are both (N,3) array
    '''
    pc2 = np.copy(pc)
    pc2[...,[0,1,2]] = pc2[...,[0,2,1]] # viz X,Y,Z = cam X,Z,-Y
    pc2[...,2] *= -1
    return pc2


def write_ply_pc(points, out_filename, num_classes=None):
    """ Color (N,3) points with RGB colors (N,3) within range [0,255] as OBJ file """
    N = points.shape[0]
    fout = open(out_filename, 'w')
    for i in range(N):
        fout.write('v %f %f %f\n' % (points[i,0],points[i,1],points[i,2]))
    fout.close()





