import cv2
import numpy as np
import pickle
import os

def fisheye_calib(camname, num):
    outputfolder = 'fisheye/'+ camname + '/'

    CHECKERBOARD = (6, 6)
    subpix_criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.1)
    calibration_flags = cv2.fisheye.CALIB_RECOMPUTE_EXTRINSIC + cv2.fisheye.CALIB_CHECK_COND + cv2.fisheye.CALIB_FIX_SKEW
    objp = np.zeros((1, CHECKERBOARD[0] * CHECKERBOARD[1], 3), np.float32)
    objp[0, :, :2] = np.mgrid[0:CHECKERBOARD[0], 0:CHECKERBOARD[1]].T.reshape(-1, 2)
    _img_shape = None
    objpoints = []  # 3d point in real world space
    imgpoints = []  # 2d points in image plane.

    if not os.path.exists(outputfolder):
        os.makedirs(outputfolder)

    # camera = cv2.VideoCapture(1)
    # i = 0
    # while (i<num):
    #     return_cal, img = camera.read()
    #     cv2.imshow("tst", img)
    #     key = cv2.waitKey(0)
    #     if key & 0xFF == 32:
    #         print(i, "done")
    #         gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    #         # Find the chess board corners
    #         ret, corners = cv2.findChessboardCorners(gray, CHECKERBOARD,
    #                                                  cv2.CALIB_CB_ADAPTIVE_THRESH + cv2.CALIB_CB_FAST_CHECK + cv2.CALIB_CB_NORMALIZE_IMAGE)
    #         # If found, add object points, image points (after refining them)
    #         if ret == True:
    #             objpoints.append(objp)
    #             cv2.cornerSubPix(gray, corners, (3, 3), (-1, -1), subpix_criteria)
    #             imgpoints.append(corners)
    #
    #         cv2.imwrite(outputfolder + str(i) + '.jpg', img)
    #         i = i+1
    # camera.release()


    for n in range(0, 40):
        img = cv2.imread(outputfolder+str(n)+'.jpg')
        if _img_shape == None:
            _img_shape = img.shape[:2]
        else:
            assert _img_shape == img.shape[:2], "All images must share the same size."
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        ret, corners = cv2.findChessboardCorners(gray, CHECKERBOARD,
              cv2.CALIB_CB_ADAPTIVE_THRESH + cv2.CALIB_CB_FAST_CHECK + cv2.CALIB_CB_NORMALIZE_IMAGE)
        if ret == True:
            objpoints.append(objp)
            cv2.cornerSubPix(gray, corners, (3, 3), (-1, -1), subpix_criteria)
            imgpoints.append(corners)

    print(len(imgpoints))
    N_OK = len(objpoints)
    K = np.zeros((3, 3))
    D = np.zeros((4, 1))
    rvecs = [np.zeros((1, 1, 3), dtype=np.float64) for i in range(N_OK)]
    tvecs = [np.zeros((1, 1, 3), dtype=np.float64) for i in range(N_OK)]
    # print(objpoints, imgpoints)
    rms, _, _, _, _ = \
        cv2.fisheye.calibrate(
            objpoints,
            imgpoints,
            gray.shape[::-1],
            K,
            D,
            rvecs,
            tvecs,
            calibration_flags,
            (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 1e-6)
        )
    print("Found " + str(N_OK) + " valid images for calibration")
    print("DIM=" + str(_img_shape[::-1]))
    print("K=np.array(" + str(K.tolist()) + ")")
    print("D=np.array(" + str(D.tolist()) + ")")

    DIM = _img_shape[::-1]
    calib_para = [K, D, DIM]

    pickle.dump(calib_para, open("fisheye/" + camname + "_calib.pkl", "wb"))

if __name__ == "__main__":
    fisheye_calib("cam3", 40)