from argparse import ArgumentParser
from pathlib import Path

import cv2
import numpy as np


def parse_pattern(value):
    try:
        cols, rows = value.lower().split("x", 1)
        return int(cols), int(rows)
    except ValueError as exc:
        raise SystemExit("pattern must use the form columnsxrows, for example 7x6") from exc


parser = ArgumentParser()
parser.add_argument("image_dir", help="Directory containing chessboard calibration images")
parser.add_argument("--pattern", default="7x6", help="Inner-corner pattern as columnsxrows")
parser.add_argument("--square-size", type=float, default=25.0, help="Square size in real units")
parser.add_argument("--output", default="calibration.npz", help="Output NumPy archive")
parser.add_argument("--min-frames", type=int, default=10, help="Minimum accepted chessboard frames")
args = parser.parse_args()

pattern_size = parse_pattern(args.pattern)
image_dir = Path(args.image_dir)
image_paths = sorted(
    path
    for path in image_dir.iterdir()
    if path.suffix.lower() in {".jpg", ".jpeg", ".png", ".tif", ".tiff"}
)
if not image_paths:
    raise SystemExit(f"no calibration images found in {image_dir}")

cols, rows = pattern_size
object_template = np.zeros((rows * cols, 3), np.float32)
object_template[:, :2] = np.mgrid[0:cols, 0:rows].T.reshape(-1, 2)
object_template *= args.square_size

criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)
object_points = []
image_points = []
image_size = None

for path in image_paths:
    image = cv2.imread(str(path))
    if image is None:
        print(f"skipped unreadable image: {path.name}")
        continue

    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    if image_size is None:
        image_size = gray.shape[::-1]
    elif image_size != gray.shape[::-1]:
        print(f"skipped different size: {path.name}")
        continue

    found, corners = cv2.findChessboardCorners(gray, pattern_size)
    if not found:
        print(f"skipped no corners: {path.name}")
        continue

    refined = cv2.cornerSubPix(gray, corners, (11, 11), (-1, -1), criteria)
    object_points.append(object_template.copy())
    image_points.append(refined)
    print(f"accepted: {path.name}")

if len(object_points) < args.min_frames:
    raise SystemExit(
        f"found {len(object_points)} usable frames, need at least {args.min_frames}"
    )

rms, camera_matrix, dist_coeffs, rvecs, tvecs = cv2.calibrateCamera(
    object_points,
    image_points,
    image_size,
    None,
    None,
)

total_error = 0.0
total_points = 0
for object_point, image_point, rvec, tvec in zip(object_points, image_points, rvecs, tvecs):
    projected, _ = cv2.projectPoints(
        object_point,
        rvec,
        tvec,
        camera_matrix,
        dist_coeffs,
    )
    error = cv2.norm(image_point, projected, cv2.NORM_L2)
    total_error += error * error
    total_points += len(object_point)

reprojection_rmse = (total_error / total_points) ** 0.5
output_path = Path(args.output)
np.savez(
    output_path,
    camera_matrix=camera_matrix,
    dist_coeffs=dist_coeffs,
    rvecs=np.asarray(rvecs),
    tvecs=np.asarray(tvecs),
    image_size=np.array(image_size),
    pattern_size=np.array(pattern_size),
    square_size=np.array([args.square_size]),
    rms=np.array([rms]),
    reprojection_rmse=np.array([reprojection_rmse]),
)

np.set_printoptions(precision=4, suppress=True)
print(f"images found: {len(image_paths)}")
print(f"accepted frames: {len(object_points)}")
print(f"image size: {image_size[0]}x{image_size[1]}")
print("camera matrix:")
print(camera_matrix)
print("distortion coefficients:")
print(dist_coeffs.ravel())
print(f"opencv rms: {rms:.4f} px")
print(f"reprojection rmse: {reprojection_rmse:.4f} px")
print(f"wrote: {output_path}")