import numpy as np
import cv2
import glob
import pprint
import re

cm1 = np.loadtxt('calibrationdata\\leftcamera\\leftcameramatrix.txt')
cm2 = np.loadtxt('calibrationdata\\rightcamera\\rightcameramatrix.txt')
dc1 = np.loadtxt('calibrationdata\\leftcamera\\leftcameradistcoeffs.txt')
dc2 = np.loadtxt('calibrationdata\\rightcamera\\rightcameradistcoeffs.txt')

# termination criteria
criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 100, 0.001)

# prepare object points, like (0,0,0), (1,0,0), (2,0,0) ....,(6,5,0)
objp = np.zeros((6*9,3), np.float32)
objp[:,:2] = np.mgrid[0:9,0:6].T.reshape(-1,2)

# Arrays to store object points and image points from all the images.
objpoints = [] # 3d point in real world space
imgpointsL = [] # 2d points in image plane.
imgpointsR = [] # 2d points in image plane.

leftImages = glob.glob('stereocalib\left\*.jpg')

for fnameL in leftImages:
    fnameR = re.sub('left', 'right', fnameL)
    imgL = cv2.imread(fnameL)
    imgR = cv2.imread(fnameR)
    height, width, depth  = imgL.shape

    grayL = cv2.cvtColor(imgL,cv2.COLOR_BGR2GRAY)
    grayR = cv2.cvtColor(imgR,cv2.COLOR_BGR2GRAY)

    # Find the chess board corners
    retL, cornersL = cv2.findChessboardCorners(grayL, (9,6),None)
    retR, cornersR = cv2.findChessboardCorners(grayR, (9,6),None)

    # If found, add object points, image points (after refining them)
    if retL == True and retR == True:
        objpoints.append(objp)

        corners2L = cv2.cornerSubPix(grayL, cornersL, (11,11), (-1,-1), criteria)
        corners2R = cv2.cornerSubPix(grayR, cornersR, (11,11), (-1,-1), criteria)
        imgpointsL.append(corners2L)
        imgpointsR.append(corners2R)

        # Draw and display the corners
        imgL = cv2.drawChessboardCorners(imgL, (9,6), corners2L, retL)
        imgR = cv2.drawChessboardCorners(imgR, (9,6), corners2R, retR)
        cv2.imshow('ORIGINAL LEFT CAMERA: ', imgL)
        cv2.imshow('ORIGINAL RIGHT CAMERA: ', imgR)
        cv2.waitKey(100)

cv2.destroyAllWindows()

# retVal = returned value, cm1 = Camera Matrix 1
# dc1 = Distortion Coefficients matrix 1, cm2 = Camera Matrix 2
# dc2 = Distortion Coefficients matrix 2
# r = rotation matrix, t = translation vector
# e = essential matrix, f = fundamental matrix
retVal, cm1, dc1, cm2, dc2, r, t, e, f = cv2.stereoCalibrate(objpoints, imgpointsL, imgpointsR, cm1, dc1, cm2, dc2, (width, height), None, None, cv2.CALIB_FIX_INTRINSIC, criteria)

print "return value"
pprint.pprint(retVal)
print "\nCamera Mat1"
pprint.pprint(cm1)
print "\nDist Coeffs1"
pprint.pprint(dc1)
print "\nCamera Mat2"
pprint.pprint(cm2)
print "\nDist Coeffs2"
pprint.pprint(dc2)
print "\nRotation"
pprint.pprint(r)
print "\nTranslation"
pprint.pprint(t)
print "\nEssential Mat"
pprint.pprint(e)
print "\nFundamental Mat"
pprint.pprint(f)

np.savetxt('calibrationdata\\cameramatrix1.txt', cm1);
np.savetxt('calibrationdata\\cameramatrix2.txt', cm2);
np.savetxt('calibrationdata\\distcoeffs1.txt', dc1);
np.savetxt('calibrationdata\\distcoeffs2.txt', dc2);
np.savetxt('calibrationdata\\rotation.txt', r);
np.savetxt('calibrationdata\\translation.txt', t);
np.savetxt('calibrationdata\\essentialmat.txt', e);
np.savetxt('calibrationdata\\fundamentalmat.txt', f);

print "Calibration done."