"""
calibrate_picam.py
------------------
Camera calibration for Raspberry Pi Camera Module using a printed
chessboard pattern (6 columns x 9 rows of squares).

Run this on the Raspberry Pi:
    python3 calibrate_picam.py

Outputs:
    camera_matrix.npy
    dist_coeffs.npy
"""

import cv2
import numpy as np
import os
import time
from picamera2 import Picamera2

# ── 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)  ← FIXED
SQUARE_SIZE   = 20.0      # mm per square (change if you used a different size)
PHOTOS_NEEDED = 20
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

# ── Start Pi Camera ────────────────────────────────────────────────
picam = Picamera2()
config = picam.create_preview_configuration(
    main={"size": (640, 480), "format": "BGR888"}
)
picam.configure(config)
picam.start()
time.sleep(2)   # camera warm-up

print("=" * 50)
print("PI CAMERA CALIBRATION")
print("=" * 50)
print(f"Chessboard inner corners expected: {CHESSBOARD[0]} x {CHESSBOARD[1]}")
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("=" * 50)

while True:
    frame = picam.capture_array()
    display = frame.copy()
    gray    = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

    found, corners = cv2.findChessboardCorners(
        gray, CHESSBOARD,
        cv2.CALIB_CB_ADAPTIVE_THRESH +
        cv2.CALIB_CB_FAST_CHECK +
        cv2.CALIB_CB_NORMALIZE_IMAGE
    )

    if found:
        corners2 = cv2.cornerSubPix(
            gray, corners, (11, 11), (-1, -1),
            (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)
        )
        cv2.drawChessboardCorners(display, CHESSBOARD, corners2, found)
        cv2.putText(display, "Chessboard FOUND - Press SPACE",
                    (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)

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

    cv2.imshow("Pi Camera Calibration - SPACE=capture, Q=quit", display)

    key = cv2.waitKey(1) & 0xFF

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

    elif key == ord(' '):
        if found:
            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

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

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

picam.stop()
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...")

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

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

    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()}")

    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_picam.py")