import cv2 as cv
import numpy as np
import glob


def find_img_points(name):
    img_points = []
    img = cv.imread(name)
    gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
    # Find the chess board corners
    ret, corners = cv.findChessboardCorners(gray, CHECKERBOARD, criteria)

    # If found, add object points, image points (after refining them)
    if ret is True:
        corners = cv.cornerSubPix(gray, corners, (11, 11), (-1, -1), criteria)
        img_points.append(corners)
        board_points = corners.astype(np.int32)

        cv.drawContours(gray, board_points, -1, 255, 5)
        cv.imshow(name, gray)
    return img_points, ret is True

assert cv.__version__[0] == '3', 'The fisheye module requires opencv version >= 3.0.0'
print(cv.__version__)
CHECKERBOARD = (6, 9)
objPoint = np.zeros((1, CHECKERBOARD[0]*CHECKERBOARD[1], 3), np.float32)
objPoint[0, :, :2] = np.mgrid[0:CHECKERBOARD[0], 0:CHECKERBOARD[1]].T.reshape(-1, 2)
criteria = (cv.TERM_CRITERIA_EPS+cv.TERM_CRITERIA_MAX_ITER, 30, 0.001)
objPoints = []  # 3d point in real world space
imgPointsLeft, imgPointsRight = [], []
imgShape = (960, 600)

images_l = glob.glob('l*.jpg')
images_r = glob.glob('r*.jpg')
if len(images_l) != len(images_r):
    raise Exception(f'len of image_left {len(images_l)} != len of image_right {len(images_r)}')
for img_lm, img_r in zip(images_l, images_r):
    l, fLeft = find_img_points(img_lm)
    r, fRight = find_img_points(img_r)

    if fLeft and fRight:
        imgPointsLeft.append(l)
        imgPointsRight.append(r)
        objPoints.append(objPoint)
N_OK = len(objPoints)
K_left = np.zeros((3, 3))
K_right = np.zeros((3, 3))
D_left = np.zeros((4, 1))
D_right = np.zeros((4, 1))

rvecs = np.zeros((1, 1, 3), dtype=np.float64)
tvecs = np.zeros((1, 1, 3), dtype=np.float64)

objPoints = np.array([objPoint]*N_OK, dtype=np.float64)
objPoints = np.reshape(objPoints, (N_OK, 1, -1, 3))

imgPointsLeft = np.asarray(imgPointsLeft, dtype=np.float64)
imgPointsRight = np.asarray(imgPointsRight, dtype=np.float64)

imgPointsLeft = np.reshape(imgPointsLeft, (N_OK, 1, -1, 2))
imgPointsRight = np.reshape(imgPointsRight, (N_OK, 1, -1, 2))

calibration_flags = cv.fisheye.CALIB_CHECK_COND+cv.fisheye.CALIB_FIX_SKEW+cv.fisheye.CALIB_RECOMPUTE_EXTRINSIC\
       + cv.CALIB_FIX_PRINCIPAL_POINT+cv.fisheye.CALIB_FIX_K1+cv.fisheye.CALIB_FIX_K2\
       + cv.fisheye.CALIB_FIX_K3+cv.fisheye.CALIB_FIX_K3

rms, K_left, D_left, K_right, D_right, rvecs, tvecs = \
cv.fisheye.stereoCalibrate(objPoints, imgPointsLeft, imgPointsRight, K_left, D_left, K_right, D_right, imgShape, rvecs,
                           tvecs, flags=calibration_flags, criteria=criteria)

maps = [cv.fisheye.initUndistortRectifyMap(K_left, D_left, np.eye(3), K_left, imgShape, cv.CV_16SC2),
        cv.fisheye.initUndistortRectifyMap(K_right, D_right, np.eye(3), K_right, imgShape, cv.CV_16SC2)]

for img_lm, img_r in zip(images_l, images_r):
    img = cv.imread(img_lm)
    undistorted_img = cv.remap(img, *maps[0], interpolation=cv.INTER_LINEAR, borderMode=cv.BORDER_CONSTANT)
    cv.imshow("undistorted_" + img_lm, undistorted_img)

    img = cv.imread(img_r)
    undistorted_img = cv.remap(img, *maps[1], interpolation=cv.INTER_LINEAR, borderMode=cv.BORDER_CONSTANT)
    cv.imshow("undistorted_" + img_r, undistorted_img)

cv.waitKey(0)
cv.destroyAllWindows()