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

class EndoCalibrator:

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

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

        resize_size = (self.args.resized_width, self.args.resized_height)
        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)
            # img = cv2.resize(img, resize_size, interpolation=cv2.INTER_LANCZOS4)
            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)
                    print(fname)
                    cv2.imshow("corners", img)
                    cv2.waitKey(0)

        ret, mtx, dist, rvecs, tvecs = cv2.calibrateCamera(obj_points, img_points, gray.shape[::-1], None, None)

        if Is_example_showed:
            for fname in chess_img_list:
                img = cv2.imread(fname)
                #img = cv2.resize(img, resize_size, interpolation=cv2.INTER_LANCZOS4)
                undistort_img = cv2.undistort(img, mtx, dist)
                cv2.imshow("undistort", undistort_img)
                cv2.waitKey()

        zoom_x = resize_size[0] / gray.shape[1]
        zoom_y = resize_size[1] / gray.shape[0]
        intrinsics_scaled = self.scale_intrinsics(mtx, zoom_x, zoom_y)
        os.makedirs(self.args.npy_outdir, exist_ok=True)
        np.save(os.path.join(self.args.npy_outdir, 'intrinsics_scaled.npy'), intrinsics_scaled)
        np.save(os.path.join(self.args.npy_outdir, 'dist_coeffs.npy'), dist)

    def scale_intrinsics(self, mat, sx, sy):
        out = np.copy(mat)
        out[0, 0] *= sx
        out[0, 2] *= sx
        out[1, 1] *= sy
        out[1, 2] *= sy
        return out