# this script is used to visualize the camera locations and the point clouds
# from easymocap

'''
  @ Date: 2022-09-26 16:32:19
  @ Author: Qing Shuai
  @ Mail: s_q@zju.edu.cn
  @ LastEditors: Qing Shuai
  @ LastEditTime: 2022-10-17 13:05:28
  @ FilePath: /EasyMocapPublic/apps/calibration/vis_camera_by_open3d.py
'''
import open3d as o3d
import os
import cv2
import numpy as np
from easymocap_tools.o3dwrapper import Vector3dVector, create_pcd
from easymocap_tools.vis_base import generate_colorbar

from calib_tools import read_json

from calib_tools import DataPath

def read_cameras(extri_path):
    extri_data = read_json(extri_path)
    cameras = {}
    for cam_name, extri in extri_data.items():
        extri = extri_data[cam_name]
        R = np.array(extri['R']).reshape(3, 3)
        T = np.array(extri['T']).reshape(3, 1)
        Rvec = None
        if 'Rvec' in extri:
            Rvec = np.array(extri['Rvec']).reshape(3, 1)
        cameras[cam_name] = {
            'R': R,
            'T': T,
            'Rvec': Rvec if Rvec is not None else ''
        }
    return cameras

# 对所有相机的旋转矩阵和平移向量进行全局变换
# 将所有相机的R T矩阵变换到新的全局坐标中
def transform_cameras(cameras):
    dims = {'x': 0, 'y': 1, 'z': 2}
    # 初始化全局变换矩阵
    R_global = np.eye(3)
    T_global = np.zeros((3, 1))
    # order: trans0, rot, trans
    # 初始化平移参数
    if len(args.trans0) == 3:
        trans = np.array(args.trans0).reshape(3, 1)
        T_global += trans
    # 处理旋转变换， args.rot = ['x', 90, 'y', 45]
    if len(args.rot) > 0:
        for i in range(len(args.rot) // 2): # 整除法
            dim = args.rot[2 * i]  # 旋转轴 ('x', 'y', 'z')
            val = float(args.rot[2 * i + 1])  # 旋转角度（度数）
            rvec = np.zeros((3,))  # 创建一个 3 维零向量，表示旋转向量
            # 构造旋转向量
            rvec[dims[dim]] = np.deg2rad(val)  # 将旋转角度转换为弧度，并填充到旋转向量中
            # 转换为旋转矩阵
            R = cv2.Rodrigues(rvec)[0]  # 使用 OpenCV 的 Rodrigues 函数将旋转向量转换为旋转矩阵
            # 累计旋转矩阵
            R_global = R @ R_global  # 将当前旋转矩阵与全局旋转矩阵累乘
        T_global = R_global @ T_global  # 最终的平移向量与旋转矩阵相乘
    # 平移相机
    if len(args.trans) == 3:
        trans = np.array(args.trans).reshape(3, 1)
        T_global += trans
    trans = np.eye(4)  # 构造全局变换矩阵
    trans[:3, :3] = R_global
    trans[:3, 3:] = T_global
    # apply the transformation of each camera
    # 对每个相机应用变换矩阵
    for key, cam in cameras.items():
        RT = np.eye(4)
        RT[:3, :3] = cam['R']
        RT[:3, 3:] = cam['T']
        RT = RT @ np.linalg.inv(trans)
        cam.pop('Rvec', '')
        cam['R'] = RT[:3, :3]
        cam['T'] = RT[:3, 3:]
    return cameras, trans


if __name__ == '__main__':
    import argparse

    parser = argparse.ArgumentParser()
    # parser.add_argument('--subs', type=str, default=[], nargs='+') # 指定需要可视化的相机的全部子集，后面删掉
    # parser.add_argument('--pcd', type=str, default=[], nargs='+') # 指定要加载的点云或三维网格文件列表。
    parser.add_argument('--trans0', type=float, nargs=3,
                        default=[], help='translation')
    parser.add_argument('--rot', type=str, nargs='+',
                        default=[], help='control the rotation')
    parser.add_argument('--trans', type=float, nargs=3,
                        default=[], help='translation')
    parser.add_argument('--debug', action='store_true')
    args = parser.parse_args()

    grids = []
    cameras = read_cameras(DataPath.extri_json_path)
    cameras, trans = transform_cameras(cameras)

    print(repr(trans))
    # for pcd in args.pcd:
    #     if not os.path.exists(pcd):
    #         print(pcd, ' not exist')
    #         continue
    #     if pcd.endswith('.npy'):
    #         data = np.load(pcd)
    #         points = data[:, :3]
    #         colors = data[:, 3:]
    #         points = (trans[:3, :3] @ points.T + trans[:3, 3:]).T
    #         p = create_pcd(points, colors=data[:, 3:])
    #         grids.append(p)
    #     elif pcd.endswith('.ply'):
    #         print('Load pcd: ', pcd)
    #         p = o3d.io.read_point_cloud(pcd)
    #         print(p)
    #         grids.append(p)
    #     elif pcd.endswith('.pkl'):
    #         p = o3d.io.read_triangle_mesh(pcd)
    #         grids.append(p)
    #     elif pcd.endswith('.obj'):
    #         p = o3d.io.read_triangle_mesh(pcd)
    #         vertices = np.asarray(p.vertices)
    #         print(vertices.shape)
    #         print(vertices.min(axis=0))
    #         print(vertices.max(axis=0))
    #         grids.append(p)

    center = o3d.geometry.TriangleMesh.create_coordinate_frame(
        size=1, origin=[0, 0, 0]) # 坐标系原点
    grids.append(center)
    colorbar = generate_colorbar(len(cameras), rand=False) # 生成颜色条
    camera_locations = []
    for ic, (cam, camera) in enumerate(cameras.items()):
        # if len(args.subs) > 0 and cam not in args.subs: continue
        center = - camera['R'].T @ camera['T'] # 相机位置
        print('{}: {}'.format(cam, center.T[0]))
        camera_locations.append(center)
        # 创建并绘制相机的坐标系
        center = o3d.geometry.TriangleMesh.create_coordinate_frame(
            size=0.5, origin=[center[0, 0], center[1, 0], center[2, 0]])
        center.rotate(camera['R'].T)
        grids.append(center)
    camera_locations = np.stack(camera_locations).reshape(-1, 3)
    o3d.visualization.draw_geometries(grids)