import numpy as np
import xmltodict
import cv2.cv2 as cv2


def load_camera_parameters(path='calibration_result.xml'):
    parameters = {'proj': {}, 'cam': {}}
    with open(path) as f:
        cam_mat = xmltodict.parse(f.read())
    # get projector and camera intrinsics
    for key in parameters:
        K_shape = int(cam_mat['opencv_storage'][f'{key}_int']['rows']), int(cam_mat['opencv_storage'][f'{key}_int']['cols'])
        parameters[key]['K'] = np.array(cam_mat['opencv_storage'][f'{key}_int']['data'].split(), dtype=float).reshape(K_shape).T
        dist_shape = int(cam_mat['opencv_storage'][f'{key}_dist']['rows']), int(cam_mat['opencv_storage'][f'{key}_dist']['cols'])
        parameters[key]['dist'] = np.array(cam_mat['opencv_storage'][f'{key}_dist']['data'].split(), dtype=float).reshape(dist_shape).T

    # get image size
    # weird casting cause the values are str(float) (eg. '123.'), but we want int
    imsize_shape = int(cam_mat['opencv_storage'][f'img_shape']['rows']), int(cam_mat['opencv_storage'][f'img_shape']['cols'])
    parameters['imsize'] = np.array([float(x) for x in cam_mat['opencv_storage']['img_shape']['data'].split()], dtype='uint16').reshape(imsize_shape).T

    # get extrinsics
    parameters['ext'] = {}
    rot_shape = int(cam_mat['opencv_storage']['rotation']['rows']), int(cam_mat['opencv_storage']['rotation']['cols'])
    parameters['ext']['R'] = np.array(cam_mat['opencv_storage'][f'rotation']['data'].split(), dtype=float).reshape(rot_shape).T
    # switched cols and rows for mult compat with R
    trans_shape = int(cam_mat['opencv_storage']['translation']['cols']), int(cam_mat['opencv_storage']['translation']['rows'])
    parameters['ext']['T'] = np.array(cam_mat['opencv_storage'][f'translation']['data'].split(), dtype=float).reshape(trans_shape).T

    return parameters


params = load_camera_parameters()

# print(params)
print(
    params['cam']['K'].shape,
    params['cam']['dist'].shape,
    params['proj']['K'].shape,
    params['proj']['dist'].shape,
    params['imsize'].shape,
    params['ext']['R'].shape,
    params['ext']['T'].shape,
)
# print(params['imsize'].reshape((2, 1)))
params['imsize'] = params['imsize'].reshape((2, 1))
# params['imsize'] = np.array([488, 688])
# print(params['imsize'].reshape((2, 1)))
# print(np.transpose(params['ext']['T'], params['ext']['T']))

R1, R2, P1, P2, Q, validPixROI1, validPixROI2 = cv2.stereoRectify(
    params['cam']['K'],
    params['cam']['dist'],
    params['proj']['K'],
    params['proj']['dist'],
    # params['imsize'],
    (688, 488),
    params['ext']['R'],
    params['ext']['T'],
)

################# SCRATCH ##############################
import math

def isclose(x, y, rtol=1.e-5, atol=1.e-8):
    return abs(x-y) <= atol + rtol * abs(y)

def euler_angles_from_rotation_matrix(R):
    '''
    From a paper by Gregory G. Slabaugh (undated),
    "Computing Euler angles from a rotation matrix
    '''
    phi = 0.0
    if isclose(R[2,0],-1.0):
        theta = math.pi/2.0
        psi = math.atan2(R[0,1],R[0,2])
    elif isclose(R[2,0],1.0):
        theta = -math.pi/2.0
        psi = math.atan2(-R[0,1],-R[0,2])
    else:
        theta = -math.asin(R[2,0])
        cos_theta = math.cos(theta)
        psi = math.atan2(R[2,1]/cos_theta, R[2,2]/cos_theta)
        phi = math.atan2(R[1,0]/cos_theta, R[0,0]/cos_theta)
    return psi, theta, phi
####################################################

# print('R1:\n', R1)
# print(euler_angles_from_rotation_matrix(R1))
# print('R2:\n', R2)
# print(euler_angles_from_rotation_matrix(R2))
# print('P1:\n', P1)
# print('P2:\n', P2)
# print('Q :\n', Q)
#
#
# print(P1.shape)

pattern = cv2.imread('kinect_pattern.png')
sampled_pattern = cv2.imread('sampled_kinect_pattern.png')
proj_rect_map1, proj_rect_map2 = cv2.initInverseRectificationMap(
# proj_rect_map1, proj_rect_map2=cv2.initUndistortRectifyMap(
    params['proj']['K'],
    params['proj']['dist'],
    R1,
    P1,
    (688, 488),
    # (1280, 800),
    cv2.CV_16SC2,
)

# print(proj_rect_map1.shape, proj_rect_map2.shape)

samp_rect_pat = cv2.remap(sampled_pattern, proj_rect_map1, proj_rect_map2, cv2.INTER_CUBIC)
rect_pat = cv2.remap(pattern, proj_rect_map1, proj_rect_map2, cv2.INTER_CUBIC)

# print(rect_pat.shape)
cv2.imshow('get rect', samp_rect_pat)
cv2.waitKey()
cv2.imshow('get rect', rect_pat)
cv2.waitKey()
# cv2.imshow(rect_pat2)
cv2.waitKey()
cv2.imwrite('rectified_sampled_pattern_new.png', samp_rect_pat)
cv2.imwrite('rectified_pattern_new.png', rect_pat)