import cv2
import numpy as np
import glob

#
# Create calibration .npz from checkerboard images
#

# Parameters
CHECKERBOARD = (10, 7)         # inner corners
square_size = 25.0 / 1000.0  # 25 mm -> meters

# Prepare object points
objp = np.zeros((CHECKERBOARD[0] * CHECKERBOARD[1], 3), np.float32)
objp[:, :2] = np.mgrid[0:CHECKERBOARD[0], 0:CHECKERBOARD[1]].T.reshape(-1, 2)
objp *= square_size

objpoints = []  # 3D points
imgpoints = []  # 2D points

# Load images
images = glob.glob("*.jpg")
print("Found images:", len(images))

img_size = None

for fname in images:
    img = cv2.imread(fname)

    print(f"Processing {fname}...")
    if img is None:
        print(f"Warning: could not read {fname}")
        continue

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

    # Save image size once
    if img_size is None:
        img_size = gray.shape[::-1]

    print(f"    Gray")
    ret, corners = cv2.findChessboardCorners(
        gray,
        CHECKERBOARD,
        flags=cv2.CALIB_CB_ADAPTIVE_THRESH +
              cv2.CALIB_CB_NORMALIZE_IMAGE +
              cv2.CALIB_CB_FAST_CHECK
    )
    print("     Corners found")

    if ret:
        corners2 = cv2.cornerSubPix(
            gray,
            corners,
            (11, 11),
            (-1, -1),
            (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 1e-6)
        )

        objpoints.append(objp)
        imgpoints.append(corners2)

        print(f"✅ Corners found in {fname}")
    else:
        print(f"❌ No corners found in {fname}")

print(f"\nTotal valid images: {len(objpoints)} / {len(images)}" )
# Sanity checks
if img_size is None:
    raise RuntimeError("No images were successfully loaded.")

if len(objpoints) == 0:
    raise RuntimeError("No chessboard corners detected in any image. Calibration failed.")

print("\n=== Sanity Checks Passed ===")
# Calibration
ret, K, D, rvecs, tvecs = cv2.calibrateCamera(
    objpoints,
    imgpoints,
    img_size,
    None,
    None
)


print("\n=== Calibration Results ===")
print("RMS reprojection error:", ret)
print("Camera matrix:\n", K)
print("Distortion coefficients:\n", D)

# Save calibration
np.savez(
    "calibration.npz",
    camera_matrix=K,
    dist_coeffs=D
)

print("\n✅ Calibration saved to calibration_XPro2_16mm.npz")