import argparse

import numpy as np
import transforms3d as tfs
import json


def get_matrix_euler_radu(x, y, z, rx, ry, rz):
    """从欧拉角构建齐次变换矩阵"""
    '''构造旋转矩阵'''
    rmat = tfs.euler.euler2mat(
        # math.radians(rx), math.radians(ry), math.radians(rz)
        rx, ry, rz
    )

    """构造齐次变换矩阵"""
    rmat = tfs.affines.compose(
        np.squeeze(np.asarray((x, y, z))),
        rmat,
        [1, 1, 1]
    )
    return rmat


def get_matrix_quat(x, y, z, rw, rx, ry, rz):
    """从单位四元数构建齐次变换矩阵"""
    '''构造旋转矩阵'''
    q = [rw, rx, ry, rz]
    rmat = tfs.quaternions.quat2mat(q)

    """构造齐次变换矩阵"""
    rmat = tfs.affines.compose(
        np.squeeze(np.asarray((x, y, z))),
        rmat,
        [1, 1, 1]
    )
    return rmat

def calculate(mode, hand_path, camera_path, hand_eye_path):
    with open(camera_path, 'r', encoding='utf-8') as f:
        camera_all = json.load(f)

    with open(hand_path, 'r', encoding='utf-8') as f:
        hand_all = json.load(f)

    with open(hand_eye_path, 'r', encoding='utf-8') as f:
        data = json.load(f)
        T_cam_in_hand = data["T"]

    sum_ = int(len(hand_all) / 7)
    T_list = []
    for i in range(sum_):
        hand = [hand_all[i*7], hand_all[i*7+1], hand_all[i*7+2], hand_all[i*7+3], hand_all[i*7+4], hand_all[i*7+5], hand_all[i*7+6]]
        camera = [camera_all[i*7], camera_all[i*7+1], camera_all[i*7+2], camera_all[i*7+3], camera_all[i*7+4], camera_all[i*7+5], camera_all[i*7+6]]
        T_hand_in_base = get_matrix_quat(hand[0], hand[1], hand[2], hand[3], hand[4], hand[5], hand[6])
        T_cal_in_camera = get_matrix_quat(camera[0], camera[1], camera[2], camera[3], camera[4], camera[5], camera[6])

        if mode == "eye_to_hand":
            T = np.linalg.inv(T_hand_in_base) @ T_cam_in_hand @ T_cal_in_camera
        else:
            T = T_hand_in_base @ T_cam_in_hand @ T_cal_in_camera
        T_list.append(T)
        print(T)

    sum_sub = np.zeros_like(T_list[0])
    for i in range(len(T_list)):
        for j in range(i + 1, len(T_list)):
            d = np.subtract(T_list[i], T_list[j])
            sum_sub = np.add(sum_sub, d)

    len_sub = len(T_list) * (len(T_list) - 1) / 2

    p = np.divide(sum_sub, len_sub)
    p[0:3 ,3:4] = p[0:3 ,3:4] * 1000
    np.set_printoptions(suppress=True, floatmode='fixed', precision=8)
    print(p)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("mode", type=str)
    parser.add_argument("hand_path", type=str)
    parser.add_argument("camera_path", type=str)
    parser.add_argument("hand_eye_path", type=str)
    args = parser.parse_args()

    calculate(args.mode, args.hand_path, args.camera_path, args.hand_eye_path)


