import cv2 as cv
import numpy as np
import glob

img_shape = (640, 480)  # image resolution
images = glob.glob('calibration_imgs/*.jpg')  # create array of all image filenames
width = 10  # width of chessboard corners
height = 7  # height of chessboard corners
criteria = (cv.TERM_CRITERIA_EPS + cv.TERM_CRITERIA_MAX_ITER, 30, 0.001)  # criteria for subpixel detection

pixel = [7.4e-3, 7.4e-3]  # (x, y) physical dimension of each pixel in mm
sensor = [4.8, 3.6]  # (x, y) physical dimension of camera sensor array in mm

board_points = np.zeros((width*height, 3), np.float32)  # generate 70x3 array of type float32
board_points[:, :2] = np.mgrid[0:width, 0:height].T.reshape(-1, 2)  # generate (x, y, z) coords of points (all z = 0)

world_points = []  # world chessboard points - normalized spacing (constant regardless of frame)
image_points = []  # image/frame chessboard points - found from subpixel detection

for img in images:  # work through each image in calibration pictures array
    frame = cv.imread(img)  # pull image
    gray = cv.cvtColor(frame, cv.COLOR_BGR2GRAY)  # convert to grayscale
    corners = cv.findChessboardCorners(gray, (width, height), None)[1]  # detect chessboard corners

    if corners is not None:  # check if there are corners
        subpix = cv.cornerSubPix(gray, corners, (11, 11), (-1, -1), criteria)  # refine corner detection
        image_points.append(subpix)  # add image corners to image points array
        world_points.append(board_points)  # add new copy of world points to world points array

# run camera calibration
ret, intrinsic, distortion = cv.calibrateCamera(world_points, image_points, img_shape, None, None)[0:3]
distortion = distortion.T  # transpose distortion coefficients

fSx = intrinsic[0, 0]  # find focal length in pixels
focal_length = fSx * pixel[0]  # find focal length in mm

f = open('task2.txt', 'w')  # open file for writing results
f.write(f'focal length in pixels:\n{fSx}\n\n')  # store focal length in pixels
f.write(f'focal length in mm:\n{focal_length}\n\n')  # store focal length
f.write(f'intrinsic paramters:\n{intrinsic}\n\n')  # store intrinsic parameters
f.write(f'distortion paramters:\n{distortion}')  # store distortion parameters
f.close()  # close file

np.savez('params.npz', fl=focal_length, intr=intrinsic, dist=distortion)  # save camera and distortion parameters

cv.destroyAllWindows()  # destroy all windows generated by script
exit()  # close out script
