import cv2
import numpy as np
import os

chessCell = 0.048
cameraFocalLength = 0.014
cameraMatrixWidth = 0.036
cameraMatrixHeight = 0.024
holoFrameWidth = 1408
holoFrameHeight = 792
gridCellsX = 6
gridCellsY = 4

print('--- Finding chessboards in images ---')

cameraCorners, holoCorners = [], []

for s in os.listdir('.'):
    if s[-7:] == 'cam.png':
        image = cv2.imread(s)
        image = cv2.resize(image, (holoFrameWidth, holoFrameHeight))
        found, corners = cv2.findChessboardCorners(image, (gridCellsX - 1, gridCellsY - 1))
        if not found:
            print(s, ': failed to find chessboard corners', sep='')
            continue
        number = int(s[:s.find('_')])
        cameraCorners.append((number, np.concatenate(corners)))
    elif s[-8:] == 'holo.jpg':
        image = cv2.imread(s)
        image = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
        found, corners = cv2.findChessboardCorners(image, (gridCellsX - 1, gridCellsY - 1))
        if not found:
            print(s, ': failed to find chessboard corners', sep='')
            continue
        number = int(s[:s.find('_')])
        holoCorners.append((number, np.concatenate(corners)))

print('Using', len(cameraCorners), 'DSLR images and', len(holoCorners), 'HoloLens images')

objectPoints = [[[chessCell * j, chessCell * i, 0] for j in range(gridCellsX - 1)] for i in range(gridCellsY - 1)]
objectPoints = np.concatenate(objectPoints)

print()
print('--- Calibrating DSLR ---')

cameraFocalLengthPix = cameraFocalLength * min(holoFrameWidth / cameraMatrixWidth, holoFrameHeight / cameraMatrixHeight)
cameraMatrix = np.array([
    [cameraFocalLengthPix, 0, holoFrameWidth / 2],
    [0, cameraFocalLengthPix, holoFrameHeight / 2],
    [0, 0, 1]
    ])

cameraCalib = cv2.calibrateCamera([objectPoints.astype('float32')] * len(cameraCorners), [s[1].astype('float32') for s in cameraCorners], (holoFrameWidth, holoFrameHeight), cameraMatrix, None, flags=cv2.CALIB_USE_INTRINSIC_GUESS)

print('Camera RMS:', cameraCalib[0])
print('Camera matrix:')
print(cameraCalib[1])
print('Camera distortion coefficients:')
print(cameraCalib[2])

cameraValues = cv2.calibrationMatrixValues(cameraCalib[1], (holoFrameWidth, holoFrameHeight), 0, 0)

print('Camera field of view:', cameraValues[:2])

print()
print('--- Calibrating HoloLens ---')

holoCalib = cv2.calibrateCamera([objectPoints.astype('float32')] * len(holoCorners), [s[1].astype('float32') for s in holoCorners], (holoFrameWidth, holoFrameHeight), None, None)

print('Holo RMS:', holoCalib[0])
print('Holo matrix:')
print(holoCalib[1])
print('Holo distortion coefficients:')
print(holoCalib[2])

holoValues = cv2.calibrationMatrixValues(holoCalib[1], (holoFrameWidth, holoFrameHeight), 0, 0)

print('HoloLens field of view:', holoValues[:2])

print()
print('--- Stereo calibrating ---')

cameraCorners = {number: corners for number, corners in cameraCorners}
holoCorners = {number: corners for number, corners in holoCorners}

commonIndices = list(set(cameraCorners) & set(holoCorners))

print('Using', len(commonIndices), 'common images')

stereoCalib = cv2.stereoCalibrate([objectPoints.astype('float32')] * len(commonIndices), [cameraCorners[i].astype('float32') for i in commonIndices], [holoCorners[i].astype('float32') for i in commonIndices],
    cameraCalib[1].astype('float32'), cameraCalib[2].astype('float32'), holoCalib[1].astype('float32'), holoCalib[2].astype('float32'), (holoFrameWidth, holoFrameHeight), flags=cv2.CALIB_FIX_INTRINSIC)

print('Stereo RMS:', stereoCalib[0])
print('Translation:', stereoCalib[6].flatten())
print('Rotation in degrees:', cv2.Rodrigues(stereoCalib[5])[0].flatten() * 180 / np.pi)
print('Rotation matrix:')
print(stereoCalib[5])