import argparse
import open3d
import numpy as np
import torch

box_colormap = [
    [1, 1, 1],
    [0, 1, 0],
    [0, 1, 1],
    [1, 1, 0],
    [1, 0, 0],
    [0, 0, 1]
]

# def parse_config():
#     parser = argparse.ArgumentParser(description='arg parser')
#     args = parser.parse_args()
#     return args, cfg

def draw_scenes(points, gt_boxes=None, ref_boxes=None, ref_labels=None, ref_scores=None, point_colors=None, draw_origin=True):
    if isinstance(points, torch.Tensor):
        points = points.cpu().numpy()
    if isinstance(gt_boxes, torch.Tensor):
        gt_boxes = gt_boxes.cpu().numpy()
    if isinstance(ref_boxes, torch.Tensor):
        ref_boxes = ref_boxes.cpu().numpy()

    vis = open3d.visualization.Visualizer()
    vis.create_window()
    # 读取viewpoint参数
    param = open3d.io.read_pinhole_camera_parameters("/home/huangruixin/Downloads/script/show_benewake_oneFrame.json")
    ctr = vis.get_view_control()

    vis.get_render_option().point_size = 1.0
    vis.get_render_option().background_color = np.zeros(3)

    # draw origin
    if draw_origin:
        axis_pcd = open3d.geometry.TriangleMesh.create_coordinate_frame(size=1.0, origin=[0, 0, 0])
        vis.add_geometry(axis_pcd)

    pts = open3d.geometry.PointCloud()
    pts.points = open3d.utility.Vector3dVector(points[:, :3])

    vis.add_geometry(pts)
    if point_colors is None:
        pts.colors = open3d.utility.Vector3dVector(np.ones((points.shape[0], 3)))
    else:
        pts.colors = open3d.utility.Vector3dVector(point_colors)

    if gt_boxes is not None:
        vis = draw_box(vis, gt_boxes, (0, 0, 1))

    if ref_boxes is not None:
        vis = draw_box(vis, ref_boxes, (0, 1, 0), ref_labels, ref_scores)

    # 转换视角
    ctr.convert_from_pinhole_camera_parameters(param)
    vis.run()
    # # 存储固定视角参数
    # param = vis.get_view_control().convert_to_pinhole_camera_parameters()
    # open3d.io.write_pinhole_camera_parameters("~/Downloads/script/show_benewake_oneFrame.json", param)
    vis.destroy_window()

def translate_boxes_to_open3d_instance(gt_boxes):
    """
             4-------- 6
           /|         /|
          5 -------- 3 .
          | |        | |
          . 7 -------- 1
          |/         |/
          2 -------- 0
    """
    center = gt_boxes[0:3]
    lwh = gt_boxes[3:6]
    axis_angles = np.array([0, 0, gt_boxes[6] + 1e-10])
    rot = open3d.geometry.get_rotation_matrix_from_axis_angle(axis_angles)
    box3d = open3d.geometry.OrientedBoundingBox(center, rot, lwh)

    line_set = open3d.geometry.LineSet.create_from_oriented_bounding_box(box3d)
    lines = np.asarray(line_set.lines)
    lines = np.concatenate([lines, np.array([[1, 4], [7, 6]])], axis=0)

    line_set.lines = open3d.utility.Vector2iVector(lines)

    return line_set, box3d

def draw_box(vis, gt_boxes, color=(0, 1, 0), ref_labels=None, score=None):
    for i in range(gt_boxes.shape[0]):
        line_set, box3d = translate_boxes_to_open3d_instance(gt_boxes[i])
        if ref_labels is None:
            line_set.paint_uniform_color(color)
        else:
            line_set.paint_uniform_color(box_colormap[ref_labels[i]])

        vis.add_geometry(line_set)

        # if score is not None:
        #     corners = box3d.get_box_points()
        #     vis.add_3d_label(corners[5], '%.2f' % score[i])
    return vis


def main():
    # args, cfg = parse_config()
    # logger = common_utils.create_logger()
    # logger.info('-----------------Quick Demo of OpenPCDet-------------------------')
    which_frame = 50
    points = np.fromfile(('/home/huangruixin/Downloads/script/benewake_overpass/%.6d.bin' % which_frame), dtype= np.float32).reshape(-1,4)
    # points = np.fromfile('/home/huangruixin/Downloads/test_bin/000001.bin', dtype= np.float32).reshape(-1,4)

    # benewake: 
    limit_range = [0, -46.08, -1, 92.16, 46.08, 4.6]
    # tanway
    # limit_range = [0, -39.68, -3, 69.12, 39.68, 1]

    mask = (points[:, 0] >= limit_range[0]) & (points[:, 0] <= limit_range[3]) \
           & (points[:, 1] >= limit_range[1]) & (points[:, 1] <= limit_range[4])

    points = points[mask]

    print(points.shape)

    pred_labels = []
    pred_boxes = []
    pred_scores = []

    with open(f'/home/huangruixin/Downloads/script/road_bst/mvlidarnet_bw/{which_frame}_cls.txt', 'r') as f:
    # with open(f'/home/huangruixin/Downloads/script/road_bst/mvlidarnet_kitti/{which_frame}_cls.txt', 'r') as f:
        for line in f.readlines():
            result = line.strip().split(' ')
            pred_labels.append(result[0])
            pred_boxes.append(result[1:8])
            pred_scores.append(result[8])

    pred_labels = np.array(pred_labels, dtype=np.int)        
    pred_boxes = np.array(pred_boxes, dtype=np.float32)
    pred_scores = np.array(pred_scores, dtype=np.float32)    

    draw_scenes(
        points=points[:, :3], ref_boxes=pred_boxes,
        ref_scores=pred_scores, ref_labels=pred_labels
    )

    # if not OPEN3D_FLAG:
    #     mlab.show(stop=True)

    # logger.info('Demo done.')


if __name__ == '__main__':
    main()