import cv2
import numpy as np
from pathlib import Path



def save_coefficients(mtx, dist, path):
    """Save camera matrix and distortion coefficients to file."""
    cv_file = cv2.FileStorage(path, cv2.FILE_STORAGE_WRITE)
    cv_file.write('K', mtx)
    cv_file.write('D', dist)
    cv_file.release()


def load_coefficients(path):
    """Load camera matrix and distortion coefficients from file."""
    cv_file = cv2.FileStorage(path, cv2.FILE_STORAGE_READ)
    camera_matrix = cv_file.getNode('K').mat()
    dist_matrix = cv_file.getNode('D').mat()
    cv_file.release()
    return [camera_matrix, dist_matrix]


def generate_matrix(filename, board_vert, board_horz):
    """Main function to calibrate camera and undistort image."""

    filename_stem = Path(filename).stem

    # Define the checkerboard pattern size and criteria for corner detection
    CHECKERBOARD = (int(board_vert)-1, int(board_horz)-1)
    #CHECKERBOARD = (9, 7)
    criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)

    # Initialize object points and image points arrays
    object_points = []  # 3D points in real world space
    image_points = []  # 2D points in image plane

    # Create the object points for the chessboard corners
    objectp3d = np.zeros((1, CHECKERBOARD[0] * CHECKERBOARD[1], 3), np.float32)
    objectp3d[0, :, :2] = np.mgrid[0:CHECKERBOARD[0], 0:CHECKERBOARD[1]].T.reshape(-1, 2)

    # Load the image and convert to grayscale
    image = cv2.imread(filename)
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

    # Find chessboard corners
    ret, corners = cv2.findChessboardCorners(gray, CHECKERBOARD, cv2.CALIB_CB_ADAPTIVE_THRESH
                                             + cv2.CALIB_CB_FAST_CHECK + cv2.CALIB_CB_NORMALIZE_IMAGE)

    # If corners found, add object points and image points to arrays
    if ret:
        object_points.append(objectp3d)
        corners2 = cv2.cornerSubPix(gray, corners, (11, 11), (-1, -1), criteria)
        image_points.append(corners2)
        image = cv2.drawChessboardCorners(image, CHECKERBOARD, corners2, ret)

    # Calibrate camera using object points and image points
    h, w = gray.shape[:2]
    try:
        ret, mtx, dist, rvecs, tvecs = cv2.calibrateCamera(object_points, image_points, gray.shape[::-1], None, None)
        # Save camera matrix and distortion coefficients to file
        save_coefficients(mtx, dist, f"{filename_stem}.yml")
        return f"{filename_stem}.yml"
    except:
        print("Please check the Chessboard Dimensions")