#!/usr/bin/env python
# coding: utf-8
import transforms3d as tfs
import numpy as np
import math
import json

import rclpy
from rclpy.node import Node


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 get_matrix_eular_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_rotvector(x, y, z, rx, ry, rz):
    """从旋转向量构建齐次变换矩阵"""
    '''构造旋转矩阵'''
    rvec = np.array([rx, ry, rz])
    theta = np.linalg.norm(rvec)
    if theta < 1e-8:
        rmat = np.eye(3)  # 小角度直接返回单位矩阵
    else:
        axis = rvec / theta
        rmat = tfs.axangles.axangle2mat(axis, theta)

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


def skew(v):
    return np.array([[0, -v[2], v[1]],
                     [v[2], 0, -v[0]],
                     [-v[1], v[0], 0]])


def R2P(T):
    """旋转矩阵 --> 修正罗德里格斯向量"""
    axis, angle= tfs.axangles.mat2axangle(T[0:3, 0:3])
    P = 2 * math.sin(angle / 2) * axis
    return P


def P2R(P):
    """修正罗德里格斯向量 --> 旋转矩阵"""
    P2 = np.dot(P.T, P)
    part_1 = (1 - P2 / 2) * np.eye(3)
    part_2 = np.add(np.dot(P, P.T), np.sqrt(4- P2) * skew(P.flatten().T))
    R = np.add(part_1, np.divide(part_2, 2))
    return R


def calculate(Hgs, Hcs):
    """计算标定矩阵"""
    # H代表矩阵、h代表标量
    Hgijs = []
    Hcijs = []
    A = []
    B = []
    size = 0
    for i in range(len(Hgs)):
        for j in range(i + 1, len(Hgs)):
            size += 1
            Hgij = np.dot(np.linalg.inv(Hgs[j]), Hgs[i])
            Hgijs.append(Hgij)
            Pgij = np.dot(2, R2P(Hgij))

            Hcij = np.dot(Hcs[j], np.linalg.inv(Hcs[i]))
            Hcijs.append(Hcij)
            Pcij = np.dot(2, R2P(Hcij))

            A.append(skew(np.add(Pgij, Pcij)))
            B.append(np.subtract(Pcij, Pgij).reshape(3, 1))

    MA = np.vstack(A)
    MB = np.vstack(B)
    Pcg_ = np.dot(np.linalg.pinv(MA), MB)
    pcg = np.sqrt(np.add(1, np.dot(Pcg_.T, Pcg_)))
    Pcg = np.dot(np.dot(2, Pcg_), np.linalg.inv(pcg))
    Rcg = P2R(Pcg).reshape(3, 3)

    A = []
    B = []
    id = 0
    for i in range(len(Hgs)):
        for j in range(i + 1, len(Hgs)):
            Hgij = Hgijs[id]
            Hcij = Hcijs[id]
            A.append(np.subtract(Hgij[0:3, 0:3], np.eye(3, 3)))
            B.append(np.subtract(np.dot(Rcg, Hcij[0:3, 3:4]), Hgij[0:3, 3:4]))
            id += 1

    MA = np.asarray(A).reshape(size * 3, 3)
    MB = np.asarray(B).reshape(size * 3, 1)
    Tcg = np.dot(np.linalg.pinv(MA), MB).reshape(3, )

    return tfs.affines.compose(Tcg, np.squeeze(Rcg), [1, 1, 1])


class Calibration(Node):
    def __init__(self, name):
        super(Calibration, self).__init__(name)
        self.sync_subscriber = None
        self.sub_camera_pose = None
        self.sub_hand_pose = None
        self.Hgs, self.Hcs = [], []
        self.hand, self.camera = [], []
        self.calibration_matrix = None

        self.declare_parameter('matrix_name', 'eye_to_hand')
        self.matrix_name = self.get_parameter('matrix_name').value

        self.declare_parameter('mode', 'eye_to_hand')
        self.mode = self.get_parameter('mode').value.lower()

        if self.mode not in ['eye_in_hand', 'eye_to_hand']:
            raise ValueError("mode must be 'eye_in_hand' or 'eye_to_hand'")

        self.declare_parameter('input', 'quat')
        self.input = self.get_parameter('input').value.lower()
        if self.input == 'eular':
            self.function = get_matrix_eular_radu
        elif self.input == 'rotvertor':
            self.function = get_matrix_rotvector
        elif self.input == 'quat':
            self.function = get_matrix_quat
        else:
            raise ValueError("INPUT ERROR")
        
        self.declare_parameter('camera_pose_path', 'camera_pose_data.json')
        self.declare_parameter('hand_pose_path', 'hand_pose_data.json')
        self.camera_pose_path = self.get_parameter('camera_pose_path').value
        self.hand_pose_path = self.get_parameter('hand_pose_path').value

        self.get_pose()
        
        self.done = False

    def get_pose(self):
        with open(f'{self.camera_pose_path}', 'r', encoding='utf-8') as f:
            self.camera = json.load(f)

        with open(f'{self.hand_pose_path}', 'r', encoding='utf-8') as f:
            self.hand = json.load(f)

        self.calculate_calibration()

        print(self.hand)
        print(self.camera)

        self.get_logger().info(f"{self.calibration_matrix}")
        hand_eye_result = {
            "T": self.calibration_matrix.tolist()
        }
        with open(f"{self.matrix_name}_matrix.json", "w") as f:
            json.dump(hand_eye_result, f, indent=4)
        self.get_logger().info(f"Save as {self.matrix_name}_matrix.json")

        self.done = True

    def calculate_data(self):
        if self.input == 'quat':
            for i in range(0, len(self.hand), 7):
                self.Hgs.append(
                    np.linalg.inv(
                        self.function(
                            self.hand[i], self.hand[i + 1], self.hand[i + 2],
                            self.hand[i + 3], self.hand[i + 4], self.hand[i + 5], self.hand[i + 6]
                        )
                    )
                    if self.mode == 'eye_to_hand' else
                    self.function(
                        self.hand[i], self.hand[i + 1], self.hand[i + 2],
                        self.hand[i + 3], self.hand[i + 4], self.hand[i + 5], self.hand[i + 6]
                    )
                )
                self.Hcs.append(
                    self.function(
                        self.camera[i], self.camera[i + 1], self.camera[i + 2],
                        self.camera[i + 3], self.camera[i + 4], self.camera[i + 5], self.camera[i + 6]
                    )
                )

        else:
            for i in range(0, len(self.hand), 6):
                self.Hgs.append(
                    np.linalg.inv(
                        self.function(
                            self.hand[i], self.hand[i + 1], self.hand[i + 2],
                            self.hand[i + 3], self.hand[i + 4], self.hand[i + 5]
                        )
                    )
                    if self.mode == 'eye_to_hand' else
                    self.function(
                        self.hand[i], self.hand[i + 1], self.hand[i + 2],
                        self.hand[i + 3], self.hand[i + 4], self.hand[i + 5]
                    )
                )
                self.Hcs.append(
                    self.function(
                        self.camera[i], self.camera[i + 1], self.camera[i + 2],
                        self.camera[i + 3], self.camera[i + 4], self.camera[i + 5]
                    )
                )


    def calculate_calibration(self):
        self.calculate_data()
        self.calibration_matrix = calculate(self.Hgs, self.Hcs)


def main(args=None):
    rclpy.init(args=args)
    node = Calibration('calibration')
    try:
        while rclpy.ok() and not node.done:
            rclpy.spin_once(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        rclpy.shutdown()


if __name__ == '__main__':
    main()

