'''
Find camera intrinsicts:
camera matrix and distortion coefficients
Largely a copy from https://longervision.github.io/2017/03/16/ComputerVision/OpenCV/opencv-internal-calibration-chessboard/


Usage:

1. Set dataset variable to point to a directory containing chessboard.mp4
2. make sure CHECKERBOARD has the nr of corners in the printed board used. Use (6,9) for https://github.com/opencv/opencv/blob/4.x/doc/pattern.png
3. Scripts creates a `calibration.json` in the dataset folder

'''

from pathlib import Path
import numpy as np
import cv2
import json
import tqdm
import math

dataset = Path('hof3-cam-baumer')

# set needed detections. Use math.inf to scan the whole video
needed_detections = math.inf # 20

# Defining the dimensions of checkerboard
CHECKERBOARD = (6,9)

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


# prepare object points, like (0,0,0), (1,0,0), (2,0,0) ....,(6,5,0)
objp = np.zeros((CHECKERBOARD[0] * CHECKERBOARD[1],3), np.float32)
objp[:,:2] = np.mgrid[0:CHECKERBOARD[0],0:CHECKERBOARD[1]].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.

videofile = dataset / "chessboard7.mp4"
cap = cv2.VideoCapture(videofile)

frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
dim = {
    'width': frame_width,
    'height': frame_height,
}
found = 0

p = tqdm.tqdm()
p2 = tqdm.tqdm(total=needed_detections)

first_found=False
no_frames_for = 0

while ((found < needed_detections) if math.isfinite(needed_detections) else True):
    ret, img = cap.read() # Capture frame-by-frame
    if not ret:
        break

    p.update()
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    # Find the chess board corners
    ret, corners = cv2.findChessboardCorners(gray, CHECKERBOARD,None)

    # If found, add object points, image points (after refining them)
    if ret == True:
        if not first_found:
            first_found = True
            print(f"first at {p.n}")
        no_frames_for = 0
        objpoints.append(objp)   # Certainly, every loop objp is the same, in 3D.
        corners2 = cv2.cornerSubPix(gray,corners,(11,11),(-1,-1),criteria)
        imgpoints.append(corners2)
        p2.update()
        p2.n

        # Draw and display the corners
        img = cv2.drawChessboardCorners(img, CHECKERBOARD, corners2, ret)
        found += 1
    else:
        no_frames_for += 1
    
    if first_found and no_frames_for > 10:
        break

    cv2.imshow('img', img)
    cv2.waitKey(1)

# When everything done, release the capture
cap.release()
cv2.destroyAllWindows()

print(f"Found {found} detections")


np.savez(dataset / "chessboard7-points.npz", objpoints=objpoints, imgpoints=imgpoints)