"""
ChArUco detection confidence test (OpenCV)

What it does:
- Opens one or multiple cameras
- Detects ArUco markers and interpolated ChArUco corners
- Draws overlays
- Prints a simple "confidence" score per camera:
    markers_found, charuco_corners_found, and a normalized confidence value

Requirements:
- OpenCV built with aruco module (opencv-contrib-python)
    pip install opencv-contrib-python
"""

import cv2
import numpy as np
import time

# ----------------------------
# USER SETTINGS (match your printed board!)
# ----------------------------
CAM_IDS = [0, 2]                 # set to [0] for single camera test, or [0,2,3,...]
USE_DSHOW_ON_WINDOWS = True      # good for Windows
RESOLUTION = (3264, 2448)        # (width, height) if your cameras support it; else set None
DICT_ID = cv2.aruco.DICT_4X4_50

SQUARES_X = 6                    # number of chessboard squares in X
SQUARES_Y = 9                    # number of chessboard squares in Y
SQUARE_LEN_MM = 25.0             # square size in mm
MARKER_LEN_MM = 18.0             # marker size in mm (must be < square)

# Confidence thresholds (tune if needed)
MIN_MARKERS_OK = 4               # markers to consider "good"
MIN_CHARUCO_OK = 15              # charuco corners to consider "good"

# Display
WINDOW_SCALE = 0.5               # downscale for display if res is huge (0.5 = half size)
PRINT_EVERY_SEC = 0.5
# ----------------------------


def open_camera(cam_id: int):
    backend = cv2.CAP_DSHOW if (USE_DSHOW_ON_WINDOWS and hasattr(cv2, "CAP_DSHOW")) else 0
    cap = cv2.VideoCapture(cam_id, backend)
    if not cap.isOpened():
        return None

    if RESOLUTION is not None:
        w, h = RESOLUTION
        cap.set(cv2.CAP_PROP_FRAME_WIDTH, float(w))
        cap.set(cv2.CAP_PROP_FRAME_HEIGHT, float(h))

    return cap


def compute_confidence(markers: int, charuco: int) -> float:
    """
    Simple normalized confidence heuristic:
    - markers contribute up to MIN_MARKERS_OK
    - charuco corners contribute up to MIN_CHARUCO_OK
    """
    m = min(markers / max(MIN_MARKERS_OK, 1), 1.0)
    c = min(charuco / max(MIN_CHARUCO_OK, 1), 1.0)
    # Weighted: charuco corners matter more for calibration quality
    return 0.35 * m + 0.65 * c


def main():
    dictionary = cv2.aruco.getPredefinedDictionary(DICT_ID)
    board = cv2.aruco.CharucoBoard(
        (SQUARES_X, SQUARES_Y),
        SQUARE_LEN_MM,
        MARKER_LEN_MM,
        dictionary
    )

    # Detector parameters
    detector_params = cv2.aruco.DetectorParameters()
    # You can tweak these if detection is unstable:
    # detector_params.adaptiveThreshWinSizeMin = 3
    # detector_params.adaptiveThreshWinSizeMax = 23
    # detector_params.adaptiveThreshWinSizeStep = 10

    detector = cv2.aruco.ArucoDetector(dictionary, detector_params)

    caps = {}
    for cam_id in CAM_IDS:
        cap = open_camera(cam_id)
        if cap is None:
            print(f"[ERR] Could not open camera {cam_id}")
        else:
            caps[cam_id] = cap
            print(f"[OK] Opened camera {cam_id}")

    if not caps:
        print("No cameras opened. Exiting.")
        return

    last_print = 0.0

    print("\nControls:")
    print("  ESC = quit")
    print("  Space = print one-shot stats immediately\n")

    try:
        while True:
            frames_vis = []
            stats = {}

            for cam_id, cap in caps.items():
                ok, frame = cap.read()
                if not ok or frame is None:
                    stats[cam_id] = (0, 0, 0.0)
                    continue

                gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
                gray = 255 - gray
                # Detect markers
                corners, ids, rejected = detector.detectMarkers(gray)
                markers_found = 0 if ids is None else len(ids)

                # Draw markers
                vis = frame.copy()
                if ids is not None:
                    cv2.aruco.drawDetectedMarkers(vis, corners, ids)

                # Interpolate ChArUco corners (requires some markers found)
                charuco_found = 0
                if ids is not None and len(ids) > 0:
                    charuco_detector = cv2.aruco.CharucoDetector(board)

                    charuco_corners = None
                    charuco_ids = None

                    if ids is not None and len(ids) > 0:
                        res = charuco_detector.detectBoard(gray)

                        # OpenCV versions differ in what they return; handle both safely.
                        # Common patterns:
                        #   (charucoCorners, charucoIds, markerCorners, markerIds)
                        #   (charucoCorners, charucoIds, rejectedMarkerCandidates)
                        #   (charucoCorners, charucoIds, ...)
                        charuco_corners = res[0] if len(res) > 0 else None
                        charuco_ids = res[1] if len(res) > 1 else None
                    charuco_found = 0
                    if charuco_ids is not None:
                        charuco_found = len(charuco_ids)
                        cv2.aruco.drawDetectedCornersCharuco(vis, charuco_corners, charuco_ids)

                conf = compute_confidence(markers_found, charuco_found)
                stats[cam_id] = (markers_found, charuco_found, conf)

                # Overlay text
                h, w = vis.shape[:2]
                lines = [
                    f"Cam {cam_id}",
                    f"Markers: {markers_found}",
                    f"ChArUco corners: {charuco_found}",
                    f"Confidence: {conf:.2f}",
                    "GOOD" if (markers_found >= MIN_MARKERS_OK and charuco_found >= MIN_CHARUCO_OK) else "MOVE / LIGHT / FOCUS"
                ]

                y = 30
                for line in lines:
                    cv2.putText(vis, line, (20, y), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0) if "GOOD" in line else (0, 200, 255), 2)
                    y += 32

                # Downscale for display
                if WINDOW_SCALE != 1.0:
                    vis = cv2.resize(vis, None, fx=WINDOW_SCALE, fy=WINDOW_SCALE, interpolation=cv2.INTER_AREA)

                frames_vis.append(vis)

            # Combine displays
            if frames_vis:
                # Stack horizontally; if many cams, wrap to multiple rows
                max_per_row = 3
                rows = []
                for i in range(0, len(frames_vis), max_per_row):
                    row = frames_vis[i:i + max_per_row]
                    # pad heights
                    max_h = max(img.shape[0] for img in row)
                    padded = []
                    for img in row:
                        if img.shape[0] < max_h:
                            pad = max_h - img.shape[0]
                            img = cv2.copyMakeBorder(img, 0, pad, 0, 0, cv2.BORDER_CONSTANT, value=(0, 0, 0))
                        padded.append(img)
                    rows.append(np.hstack(padded))
                grid = np.vstack(rows)

                cv2.imshow("ChArUco Confidence Test", grid)

            key = cv2.waitKey(1) & 0xFF
            now = time.time()

            if key == 27:  # ESC
                break

            if key == 32:  # Space
                last_print = 0  # force print now

            if now - last_print >= PRINT_EVERY_SEC:
                last_print = now
                # Print compact stats
                msg = " | ".join(
                    f"cam{cid}: M={m} C={c} conf={conf:.2f}"
                    for cid, (m, c, conf) in sorted(stats.items())
                )
                print(msg)

    finally:
        for cap in caps.values():
            cap.release()
        cv2.destroyAllWindows()


if __name__ == "__main__":
    main()