import numpy as np
import cv2 as cv
import glob

PAT_NX = 7
PAT_NY = 6

def main():

	# termination criteria
	criteria = (cv.TERM_CRITERIA_EPS + cv.TERM_CRITERIA_MAX_ITER, 30, 0.01)

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

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

	images = glob.glob('*.jpg')

	for fname in images:
		img = cv.imread(fname)
		gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY)

		# Find the chess board corners
		ret, corners = cv.findChessboardCorners(gray, (PAT_NX, PAT_NY), None)

		# If found, add object points, image points (after refining them)
		if ret:
			objpoints.append(objp)

			corners2 = cv.cornerSubPix(gray, corners, (11, 11), (-1, -1), criteria)
			imgpoints.append(corners2)

			# Draw and display the corners
			cv.drawChessboardCorners(img, (PAT_NX, PAT_NY), corners2, ret)
			cv.imshow('img', img)
			cv.waitKey(-1)

		#  Calibration
		ret, mtx, dist, rvecs, tvecs = cv.calibrateCamera(objpoints, imgpoints, gray.shape[::-1], None, None)

		# write to file
		np.savez("results/cam.npz", mtx=mtx, dist=dist, rvecs=rvecs, tvecs=tvecs)

		#  Undistortion
		img = cv.imread('checker_cam.jpg')
		h, w = img.shape[:2]
		newcameramtx, roi = cv.getOptimalNewCameraMatrix(mtx, dist, (w, h), 1, (w, h))

		#  1. Undistort using cv.undistort()
		dst = cv.undistort(img, mtx, dist, None, newcameramtx)

		# crop the image
		x, y, w, h = roi
		dst = dst[y:y + h, x:x + w]
		cv.imwrite('results/calibresult_1.png', dst)

		# 2. Undistort using remapping
		mapx, mapy = cv.initUndistortRectifyMap(mtx, dist, None, newcameramtx, (w, h), 5)
		dst = cv.remap(img, mapx, mapy, cv.INTER_LINEAR)

		# crop the image
		x, y, w, h = roi
		dst = dst[y:y + h, x:x + w]
		cv.imwrite('results/calibresult_2.png', dst)

		#  Re-projection Error
		mean_error = 0
		for i in range(len(objpoints)):
			imgpoints2, _ = cv.projectPoints(objpoints[i], rvecs[i], tvecs[i], mtx, dist)
			error = cv.norm(imgpoints[i], imgpoints2, cv.NORM_L2) / len(imgpoints2)
			mean_error += error

		print("total error: {}".format(mean_error / len(objpoints)))
	cv.destroyAllWindows()

if __name__ == '__main__':
	main()