"""
calibrate.py
------------
Camera calibration for a USB webcam (laptop) using a printed
chessboard pattern.

Run this on your laptop:
    python calibrate.py

Outputs:
    camera_matrix.npy
    dist_coeffs.npy
"""

import cv2
import numpy as np
import os

# ── Settings ──────────────────────────────────────────────────────
# Your printed board: 6 columns x 9 rows of SQUARES
# Inner corners = squares - 1
# Order required by OpenCV: (corners_horizontal, corners_vertical)
CHESSBOARD    = (5, 8)    # (cols-1, rows-1) = (6-1, 9-1)
SQUARE_SIZE   = 28.0      # mm per square - MEASURE YOUR ACTUAL PRINTED SQUARE!
PHOTOS_NEEDED = 20
CAMERA_INDEX  = 0         # 0 = laptop webcam
SAVE_DIR      = "calib_photos"
# ──────────────────────────────────────────────────────────────────

os.makedirs(SAVE_DIR, exist_ok=True)

# Prepare 3D object points (same for every photo)
objp = np.zeros((CHESSBOARD[0] * CHESSBOARD[1], 3), np.float32)
objp[:, :2] = np.mgrid[0:CHESSBOARD[0], 0:CHESSBOARD[1]].T.reshape(-1, 2)
objp *= SQUARE_SIZE

objpoints = []   # 3D points in real world
imgpoints = []   # 2D points in image
photo_count = 0

cap = cv2.VideoCapture(CAMERA_INDEX)
print(f"Camera opened: {cap.isOpened()}")
print()
print("=" * 50)
print("CAMERA CALIBRATION")
print("=" * 50)
print(f"Chessboard inner corners expected: {CHESSBOARD[0]} x {CHESSBOARD[1]}")
print(f"Square size: {SQUARE_SIZE} mm")
print(f"Need {PHOTOS_NEEDED} good photos")
print()
print("Instructions:")
print("  - Keep camera FIXED, do not move it")
print("  - Hold chessboard in front of camera")
print("  - Move chessboard to different positions/angles/distances")
print("  - Press SPACE to capture when 'FOUND' shows in green")
print("  - Press Q to quit early")
print()
print("Good positions to try:")
print("  - Center, near and far")
print("  - Left, right, top, bottom of frame")
print("  - Tilted left, right, up, down")
print("  - Each corner of the frame")
print("=" * 50)

while True:
    ret, frame = cap.read()
    if not ret:
        print("Camera error!")
        break

    display = frame.copy()
    gray    = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

    # Try to find chessboard
    found, corners = cv2.findChessboardCorners(
        gray, CHESSBOARD,
        cv2.CALIB_CB_ADAPTIVE_THRESH +
        cv2.CALIB_CB_FAST_CHECK +
        cv2.CALIB_CB_NORMALIZE_IMAGE
    )

    if found:
        # Refine corners
        corners2 = cv2.cornerSubPix(
            gray, corners, (11, 11), (-1, -1),
            (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)
        )
        # Draw corners
        cv2.drawChessboardCorners(display, CHESSBOARD, corners2, found)
        cv2.putText(display, "Chessboard FOUND - Press SPACE to capture",
                    (10, 30), cv2.FONT_HERSHEY_SIMPLEX,
                    0.7, (0, 255, 0), 2)
    else:
        cv2.putText(display, "Looking for chessboard...",
                    (10, 30), cv2.FONT_HERSHEY_SIMPLEX,
                    0.7, (0, 0, 255), 2)

    # Progress counter
    cv2.putText(display,
                f"Photos captured: {photo_count} / {PHOTOS_NEEDED}",
                (10, 65), cv2.FONT_HERSHEY_SIMPLEX,
                0.7, (255, 255, 0), 2)

    cv2.imshow("Camera Calibration - Press SPACE to capture, Q to quit", display)

    key = cv2.waitKey(1) & 0xFF

    if key == ord('q'):
        print("Quitting early...")
        break

    elif key == ord(' '):
        if found:
            # Refine and save
            corners2 = cv2.cornerSubPix(
                gray, corners, (11, 11), (-1, -1),
                (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)
            )
            objpoints.append(objp)
            imgpoints.append(corners2)
            photo_count += 1

            # Save photo
            fname = os.path.join(SAVE_DIR, f"calib_{photo_count:02d}.jpg")
            cv2.imwrite(fname, frame)
            print(f"  Captured photo {photo_count}/{PHOTOS_NEEDED} -> {fname}")

            if photo_count >= PHOTOS_NEEDED:
                print(f"\nGot {PHOTOS_NEEDED} photos! Running calibration...")
                break
        else:
            print("  No chessboard found - reposition and try again")

cap.release()
cv2.destroyAllWindows()

# ── Run Calibration ───────────────────────────────────────────────
if photo_count < 5:
    print(f"\nNeed at least 5 photos, only got {photo_count}. Try again!")
else:
    print(f"\nCalibrating with {photo_count} photos...")

    h, w = cv2.imread(
        os.path.join(SAVE_DIR, "calib_01.jpg")
    ).shape[:2]

    ret, camera_matrix, dist_coeffs, rvecs, tvecs = cv2.calibrateCamera(
        objpoints, imgpoints, (w, h), None, None
    )

    # ── Results ────────────────────────────────────────────────────
    print("\n" + "=" * 50)
    print("CALIBRATION RESULTS")
    print("=" * 50)
    print(f"\nReprojection error: {ret:.4f}")
    print("(Good if below 1.0, excellent if below 0.5)")
    print(f"\nCamera Matrix:")
    print(f"  fx = {camera_matrix[0,0]:.2f} px")
    print(f"  fy = {camera_matrix[1,1]:.2f} px")
    print(f"  cx = {camera_matrix[0,2]:.2f} px")
    print(f"  cy = {camera_matrix[1,2]:.2f} px")
    print(f"\nDistortion Coefficients:")
    print(f"  {dist_coeffs.ravel()}")

    # ── Save Results ────────────────────────────────────────────────
    np.save("camera_matrix.npy",  camera_matrix)
    np.save("dist_coeffs.npy",    dist_coeffs)

    print("\nSaved:")
    print("  camera_matrix.npy")
    print("  dist_coeffs.npy")
    print("\nCalibration complete!")
    print("These files will be auto-loaded by detect_and_transform.py")