import cv2
import numpy as np
import glob
import os
import sys
import argparse

#
# Create calibration .npz from checkerboard images
# Usage:
#   python callibriate.py --camera cam0 --input-dir data/calibration/20260610_143022 [--output-dir ...]
#

# ── CLI-Parameter ──────────────────────────────────────────────────────────────
parser = argparse.ArgumentParser(description='Camera calibration from checkerboard images')
parser.add_argument('--camera',      required=True,  help='Camera ID prefix, e.g. cam0')
parser.add_argument('--input-dir',   default='.',    help='Directory containing calibration images')
parser.add_argument('--output-dir',  default=None,   help='Directory to save .npz (default: same as --input-dir)')
args = parser.parse_args()

camera     = args.camera
input_dir  = os.path.abspath(args.input_dir)
output_dir = os.path.abspath(args.output_dir) if args.output_dir else input_dir

def log(msg):
    print(msg, flush=True)   # flush=True → Node.js sieht die Zeile sofort

# ── 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 ───────────────────────────────────────────────────────────────
pattern = os.path.join(input_dir, f"{camera}_*.jpg")
images  = sorted(glob.glob(pattern))
log(f"Camera:      {camera}")
log(f"Input-Dir:   {input_dir}")
log(f"Output-Dir:  {output_dir}")
log(f"Pattern:     {pattern}")
log(f"Found images: {len(images)}")

img_size = None

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

    log(f"Processing {os.path.basename(fname)} ...")
    if img is None:
        log(f"  Warning: could not read {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,
        flags=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)
        log(f"  ✅ Corners found")
    else:
        log(f"  ❌ No corners found")

log(f"\nTotal valid images: {len(objpoints)} / {len(images)}")

# ── Sanity checks ─────────────────────────────────────────────────────────────
if img_size is None:
    log("ERROR: No images were successfully loaded.")
    sys.exit(1)

if len(objpoints) == 0:
    log("ERROR: No chessboard corners detected in any image. Calibration failed.")
    sys.exit(1)

log("\n=== Sanity Checks Passed ===")

# ── Calibration ───────────────────────────────────────────────────────────────
ret, K, D, rvecs, tvecs = cv2.calibrateCamera(
    objpoints,
    imgpoints,
    img_size,
    None,
    None
)

log("\n=== Calibration Results ===")
log(f"RMS reprojection error: {ret}")
log(f"Camera matrix:\n{K}")
log(f"Distortion coefficients:\n{D}")

# ── Save calibration ──────────────────────────────────────────────────────────
os.makedirs(output_dir, exist_ok=True)
output_path = os.path.join(output_dir, f"{camera}_calibration.npz")
np.savez(output_path, camera_matrix=K, dist_coeffs=D)

log(f"\n✅ Calibration saved to {output_path}")