import numpy as np
import cv2
import os
from glob import glob

class EndoCalibrator:

    def __init__(self, args):
        self.args = args

    def calibrate(self, Is_example_showed=True):
        criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)

        square_size = self.args.square_size
        pattern_size = (self.args.pattern_size_row, self.args.pattern_size_col)
        pattern_points = np.zeros((np.prod(pattern_size), 3), dtype=np.float32)
        pattern_points[:, :2] = np.indices(pattern_size).T.reshape(-1, 2)
        pattern_points *= square_size
        obj_points = []
        img_points = []

        chess_img_list = glob(os.path.join(self.args.imgs_dir, '*' + self.args.chess_img_extend))

        for fname in chess_img_list:
            img = cv2.imread(fname)
            gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
            ret, corners = cv2.findChessboardCorners(gray, pattern_size)

            if ret:
                obj_points.append(pattern_points)
                corners2 = cv2.cornerSubPix(gray, corners, (11, 11), (-1, -1), criteria)
                img_points.append(corners2)

                if Is_example_showed:
                    img = cv2.drawChessboardCorners(img, pattern_size, corners2, ret)
                    cv2.imshow("corners", img)
                    cv2.waitKey(0)