#!/usr/bin/env python3
import sys
sys.path.insert(0, '/home/john/opencv/build/python_loader')
import cv2 as cv
import numpy as np
import transformations as tf
from math import acos

print("OpenCV version", cv.__version__)

np.random.seed(42)
def small_random_T():
    q = np.array([1, 0, 0, 0]) + 0.01*np.random.randn(4)
    q /= np.linalg.norm(q)
    return tf.translation_matrix(0.001 * np.random.randn(3)) @ tf.quaternion_matrix(q)

def random_T():
    return tf.translation_matrix(np.random.randn(3)) @ tf.quaternion_matrix(np.random.randn(4))

num_samples = 25

T__imu__t265 = np.array([
    [  0,  0, -1,  0.1],
    [  1,  0,  0,  0.2],
    [  0, -1,  0,  0.3],
    [  0,  0,  0,  1],
    ])

T__wi__wt = random_T()
T__wt__wi = tf.inverse_matrix(T__wi__wt)

# Transforms between the IMU and its world frame (wi)
R__wi__imu = []
t__wi__imu = []

# Transforms between the T265 and its world frame (wt)
R__t265__wt = []
t__t265__wt = []

zeros = []

for use_translation, add_noise in [(True, True), (True, False), (False, True), (False, False)]:
    print("Using translation." if use_translation else "Not using translation.",
        "With noise." if add_noise else "No noise.")
    for i in range(25):
        T__wi__imu = random_T()
        T__wt__t265 = T__wt__wi @ T__wi__imu @ T__imu__t265

        if add_noise:
            T__wi__imu @= small_random_T()
            T__wt__t265 @= small_random_T()

        T__imu__wi = tf.inverse_matrix(T__wi__imu)
        T__t265__wt = tf.inverse_matrix(T__wt__t265)

        R__wi__imu.append(T__wi__imu[:3,:3])
        t__wi__imu.append(T__wi__imu[:3, 3])
        R__t265__wt.append(T__t265__wt[:3,:3])
        t__t265__wt.append(T__t265__wt[:3, 3])
        zeros.append(np.zeros((3,)))

    print("Method      rotation error   translation error")
    for method_name, method in (
            ("None", None),
            ("TSAI", cv.CALIB_HAND_EYE_TSAI),
            ("PARK", cv.CALIB_HAND_EYE_PARK),
            ("HORAUD", cv.CALIB_HAND_EYE_HORAUD),
            ("ANDREFF", cv.CALIB_HAND_EYE_ANDREFF),
            ("DANIILIDIS", cv.CALIB_HAND_EYE_DANIILIDIS),
            ):
        try: R_cal, t_cal = cv.calibrateHandEye(R__wi__imu, t__wi__imu if use_translation else zeros,
                                                R__t265__wt, t__t265__wt if use_translation else zeros,
                                                method=method)
        except:
            print(f"     {method_name} threw an exception")
            continue
        rotation_error_radians = acos((np.trace(R_cal @ T__imu__t265[:3, :3].T) - 1) / 2)
        translation_error_meters = np.linalg.norm(t_cal.squeeze() - use_translation*T__imu__t265[:3,3].squeeze())
        print(f"  {method_name:11s} {rotation_error_radians:1.6f}        {translation_error_meters: 1.6f}")
    print()