"""
Berechnet eine Kamera-Kalibrierung aus Schachbrett-Fotos und speichert das
Ergebnis als .npz.

Aufruf:
    python calibrate.py <Bildordner> <Ausgabedatei.npz>

Beispiel:
    python setup/calibrate.py test/files/cam0 data/calibration/cam0/calibration.npz

Erwartet *.jpg-Dateien im Bildordner, die ein 10x7-Eck-Schachbrettmuster
(= 11x8 Felder, 25 mm/Feld) zeigen.
"""

import sys
import glob
import pathlib
import numpy as np
import cv2

CHECKERBOARD = (10, 7)       # innere Ecken
SQUARE_MM    = 25.0 / 1000   # 25 mm in Metern

def calibrate(image_dir: str, out_path: str) -> None:
    images = sorted(glob.glob(str(pathlib.Path(image_dir) / "*.jpg")))
    if not images:
        sys.exit(f"Keine .jpg-Dateien in {image_dir!r}")
    print(f"Gefundene Bilder: {len(images)}")

    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_MM

    objpoints, imgpoints = [], []
    img_size = None

    for fname in images:
        img = cv2.imread(fname)
        if img is None:
            print(f"  WARNUNG: konnte nicht lesen: {fname}")
            continue
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        if img_size is None:
            img_size = gray.shape[::-1]

        ret, corners = cv2.findChessboardCorners(
            gray, CHECKERBOARD,
            cv2.CALIB_CB_ADAPTIVE_THRESH | cv2.CALIB_CB_NORMALIZE_IMAGE | cv2.CALIB_CB_FAST_CHECK,
        )
        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"  OK  {pathlib.Path(fname).name}")
        else:
            print(f"  --  {pathlib.Path(fname).name}  (keine Ecken gefunden)")

    print(f"\nGültige Bilder: {len(objpoints)} / {len(images)}")
    if not objpoints:
        sys.exit("Kalibrierung fehlgeschlagen: keine Ecken in keinem Bild gefunden.")

    rms, K, D, _rvecs, _tvecs = cv2.calibrateCamera(objpoints, imgpoints, img_size, None, None)
    print(f"RMS-Fehler:       {rms:.4f} px")
    print(f"Kameramatrix:\n{K}")
    print(f"Verzerrungskoeff: {D}")

    pathlib.Path(out_path).parent.mkdir(parents=True, exist_ok=True)
    np.savez(out_path, camera_matrix=K, dist_coeffs=D)
    print(f"\nGespeichert: {out_path}")

if __name__ == "__main__":
    if len(sys.argv) != 3:
        sys.exit(__doc__)
    calibrate(sys.argv[1], sys.argv[2])