"""Multi-frame PINHOLE intrinsics calibration for a named camera (scene / wrist).

Same pinhole model as the rest of the pipeline (calib_utils.estimate_focal): principal
point fixed at the image center, square pixels, ZERO distortion — we solve only the focal
length. The only change from the old single-image estimate is that we feed it ~60 board
views instead of one (a proper multi-view solve). No fisheye / distortion / pp offset.

Hold the arm-base ArUco board in front of the camera and move it around. It auto-captures
frames where the board is seen, then runs the pinhole solve and saves intrinsics/<name>.json
+ registers the camera in cameras.json.

Usage:
  python calibrate_camera.py <name> [--index N] [--board arm_base] [--frames 60]

Keys while collecting: 'c' capture now, 'q' finish early (>=8 frames), ESC abort.
"""
import argparse
import sys
import time

import numpy as np
import cv2

from aruco_boards import DICT, ARM_BASE, UMI_V2
from calib_board import CALIB_BOARD
from calib_utils import detect_boards, estimate_focal
from cameras import open_index, register, save_intrinsics

BOARDS = {"arm_base": ARM_BASE, "umi_v2": UMI_V2, "calib": CALIB_BOARD}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("name", help="camera role to save under (e.g. scene / wrist)")
    ap.add_argument("--index", type=int, default=0)
    ap.add_argument("--board", choices=list(BOARDS), default="calib")
    ap.add_argument("--frames", type=int, default=60)
    a = ap.parse_args()

    spec = BOARDS[a.board]
    name = spec["name"]
    detector = cv2.aruco.ArucoDetector(DICT, cv2.aruco.DetectorParameters())

    cap, fr = open_index(a.index)
    if cap is None or fr is None:
        sys.exit(f"Camera {a.index} did not open / delivered no frames.")
    Wn, Hn = fr.shape[1], fr.shape[0]
    print(f"Calibrating '{a.name}' on camera {a.index} ({Wn}x{Hn}) with the {a.board} board.")
    print("Move the board around to vary pose/position. 'c' capture, 'q' finish, ESC abort.\n")

    views = []          # list of (objpts Nx3, imgpts Nx2), one per captured frame
    last_t = 0.0
    while len(views) < a.frames:
        ok, frame = cap.read()
        if not ok or frame is None:
            continue
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        dets = detect_boards(gray, [spec])         # same detection the solver uses
        corners, ids, _ = detector.detectMarkers(gray)
        n = 0 if ids is None else len(ids)
        vis = frame.copy()
        if n:
            cv2.aruco.drawDetectedMarkers(vis, corners, ids)

        key = cv2.waitKey(1) & 0xFF
        now = time.time()
        if name in dets and (key == ord("c") or now - last_t > 0.4):
            views.append(dets[name])
            last_t = now

        ok_board = name in dets
        color = (0, 255, 0) if ok_board else (0, 165, 255)
        cv2.putText(vis, f"{a.name}: captured {len(views)}/{a.frames}   markers {n}",
                    (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, color, 2)
        cv2.imshow("calibrate camera (c/q/ESC)", vis)
        if key == ord("q") and len(views) >= 8:
            break
        if key == 27:
            cap.release()
            cv2.destroyAllWindows()
            sys.exit("Aborted.")

    cap.release()
    cv2.destroyAllWindows()
    if len(views) < 8:
        sys.exit(f"Only {len(views)} frames — need at least 8. Try again with more board motion.")

    print(f"\nPinhole solve over {len(views)} frames...")
    K, _, err = estimate_focal(views, (Wn, Hn), f_guess=0.9 * max(Wn, Hn))
    dist = np.zeros(5)                              # pinhole: no distortion
    print(f"  reprojection error: {err:.3f} px")
    print(f"  focal f = {K[0,0]:.1f}px   (fovy {2*np.degrees(np.arctan(Hn/2/K[0,0])):.1f} deg), "
          f"principal point fixed at ({K[0,2]:.0f}, {K[1,2]:.0f})")

    register(a.name, a.index, Wn, Hn)
    save_intrinsics(a.name, K, dist, Wn, Hn, err, len(views))
    print(f"\nSaved intrinsics/{a.name}.json (pinhole, zero distortion) and registered '{a.name}' in cameras.json.")


if __name__ == "__main__":
    main()