#!/usr/bin/env python3
"""
Detect calibration board in images, estimate camera pose via solvePnP,
and output quaternions (qx,qy,qz,qw) for up to N images.

Supports:
- chessboard or Charuco detection
- YAML/JSON camera intrinsics (OpenCV-style keys)
- Save to CSV and print to stdout

Example:
  python src/scripts/images_to_quaternions.py \
    --images data/images/*.png \
    --camera-calib data/camera.yaml \
    --pattern chessboard --board-rows 6 --board-cols 9 --square-size 0.024 \
    --limit 12 --out /tmp/quats.csv
"""
from __future__ import annotations

import argparse
import glob
import json
import os
from typing import Tuple, List

import cv2
import numpy as np


def load_camera_calib(path: str) -> Tuple[np.ndarray, np.ndarray]:
    import yaml
    with open(path, 'r') as f:
        txt = f.read()
    try:
        data = yaml.safe_load(txt)
    except Exception:
        data = json.loads(txt)
    K = data.get('camera_matrix') or data.get('cameraMatrix') or data.get('K')
    if isinstance(K, dict) and 'data' in K:
        K = K['data']
    D = data.get('dist_coeffs') or data.get('distCoeffs') or data.get('D')
    if isinstance(D, dict) and 'data' in D:
        D = D['data']
    if K is None or D is None:
        raise ValueError('Calibration must contain camera_matrix/cameraMatrix/K and dist_coeffs/distCoeffs/D')
    K = np.array(K, dtype=float).reshape(3, 3)
    D = np.array(D, dtype=float).reshape(-1, 1)
    return K, D


def build_chessboard_object_points(rows: int, cols: int, square: float) -> np.ndarray:
    objp = np.zeros((rows * cols, 3), np.float32)
    grid = np.mgrid[0:cols, 0:rows].T.reshape(-1, 2)
    objp[:, :2] = grid * square
    return objp


def detect_chessboard(gray: np.ndarray, rows: int, cols: int) -> Tuple[bool, np.ndarray]:
    flags = cv2.CALIB_CB_ADAPTIVE_THRESH + cv2.CALIB_CB_NORMALIZE_IMAGE
    ok, corners = cv2.findChessboardCorners(gray, (cols, rows), flags)
    if not ok:
        return False, None
    criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)
    corners = cv2.cornerSubPix(gray, corners, (11, 11), (-1, -1), criteria)
    return True, corners.reshape(-1, 2)


def detect_charuco(gray: np.ndarray, rows: int, cols: int, square: float, marker: float) -> Tuple[bool, np.ndarray, np.ndarray]:
    if not hasattr(cv2, 'aruco'):
        raise RuntimeError('opencv-contrib-python is required for Charuco detection')
    aruco = cv2.aruco
    aruco_dict = aruco.getPredefinedDictionary(getattr(aruco, 'DICT_4X4_50'))
    board = aruco.CharucoBoard((cols, rows), square, marker, aruco_dict)
    params = aruco.DetectorParameters()
    detector = aruco.ArucoDetector(aruco_dict, params)
    corners, ids, _ = detector.detectMarkers(gray)
    if ids is None or len(ids) == 0:
        return False, None, None
    retval, charuco_corners, charuco_ids = aruco.interpolateCornersCharuco(corners, ids, gray, board)
    ok = retval is not None and retval >= 4
    return ok, (charuco_corners.reshape(-1, 2) if ok else None), (charuco_ids if ok else None)


def rot_to_quat(R: np.ndarray) -> np.ndarray:
    tr = np.trace(R)
    if tr > 0.0:
        S = np.sqrt(tr + 1.0) * 2.0
        qw = 0.25 * S
        qx = (R[2, 1] - R[1, 2]) / S
        qy = (R[0, 2] - R[2, 0]) / S
        qz = (R[1, 0] - R[0, 1]) / S
    else:
        if R[0, 0] > R[1, 1] and R[0, 0] > R[2, 2]:
            S = np.sqrt(1.0 + R[0, 0] - R[1, 1] - R[2, 2]) * 2.0
            qx = 0.25 * S
            qy = (R[0, 1] + R[1, 0]) / S
            qz = (R[0, 2] + R[2, 0]) / S
            qw = (R[2, 1] - R[1, 2]) / S
        elif R[1, 1] > R[2, 2]:
            S = np.sqrt(1.0 + R[1, 1] - R[0, 0] - R[2, 2]) * 2.0
            qx = (R[0, 1] + R[1, 0]) / S
            qy = 0.25 * S
            qz = (R[1, 2] + R[2, 1]) / S
            qw = (R[0, 2] - R[2, 0]) / S
        else:
            S = np.sqrt(1.0 + R[2, 2] - R[0, 0] - R[1, 1]) * 2.0
            qx = (R[0, 2] + R[2, 0]) / S
            qy = (R[1, 2] + R[2, 1]) / S
            qz = 0.25 * S
            qw = (R[1, 0] - R[0, 1]) / S
    q = np.array([qx, qy, qz, qw])
    return q / (np.linalg.norm(q) + 1e-12)


def pnp_quaternion_for_image(image_path: str, K: np.ndarray, D: np.ndarray,
                              pattern: str, rows: int, cols: int, square: float,
                              charuco_marker: float | None = None) -> Tuple[np.ndarray, np.ndarray]:
    gray = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
    if gray is None:
        raise FileNotFoundError(f'Cannot read image: {image_path}')
    if pattern == 'chessboard':
        ok, corners = detect_chessboard(gray, rows, cols)
        if not ok:
            raise RuntimeError(f'Chessboard detection failed: {image_path}')
        objp = build_chessboard_object_points(rows, cols, square)
        retval, rvec, tvec = cv2.solvePnP(objp, corners, K, D, flags=cv2.SOLVEPNP_ITERATIVE)
        if not retval:
            raise RuntimeError(f'solvePnP failed: {image_path}')
    elif pattern == 'charuco':
        if charuco_marker is None:
            raise ValueError('charuco requires --charuco-marker')
        ok, charuco_corners, charuco_ids = detect_charuco(gray, rows, cols, square, charuco_marker)
        if not ok:
            raise RuntimeError(f'Charuco detection failed: {image_path}')
        aruco = cv2.aruco
        aruco_dict = aruco.getPredefinedDictionary(getattr(aruco, 'DICT_4X4_50'))
        board = aruco.CharucoBoard((cols, rows), square, charuco_marker, aruco_dict)
        obj_pts = []
        img_pts = []
        for idx, corner in zip(charuco_ids.flatten(), charuco_corners.reshape(-1, 2)):
            obj_pts.append(board.getChessboardCorners()[idx][0])
            img_pts.append(corner)
        obj_pts = np.array(obj_pts, dtype=np.float32)
        img_pts = np.array(img_pts, dtype=np.float32)
        retval, rvec, tvec = cv2.solvePnP(obj_pts, img_pts, K, D, flags=cv2.SOLVEPNP_ITERATIVE)
        if not retval:
            raise RuntimeError(f'solvePnP failed: {image_path}')
    else:
        raise ValueError('pattern must be chessboard or charuco')

    R, _ = cv2.Rodrigues(rvec)
    q = rot_to_quat(R)
    return q, tvec.flatten()


def main():
    parser = argparse.ArgumentParser(description='Estimate quaternions from board detections in images (solvePnP).')
    parser.add_argument('--images', type=str, required=True, help='Directory or glob for images')
    parser.add_argument('--camera-calib', type=str, required=True, help='Camera intrinsics (YAML/JSON)')
    parser.add_argument('--pattern', type=str, default='chessboard', choices=['chessboard', 'charuco'])
    parser.add_argument('--board-rows', type=int, default=6)
    parser.add_argument('--board-cols', type=int, default=9)
    parser.add_argument('--square-size', type=float, default=0.024, help='meters')
    parser.add_argument('--charuco-marker', type=float, default=None, help='meters for Charuco marker size')
    parser.add_argument('--limit', type=int, default=12, help='Max number of images to process')
    parser.add_argument('--out', type=str, default=None, help='Output CSV path')

    args = parser.parse_args()

    # Collect images
    if os.path.isdir(args.images):
        img_paths = sorted(glob.glob(os.path.join(args.images, '*')))
    else:
        img_paths = sorted(glob.glob(args.images))
    if not img_paths:
        raise FileNotFoundError(f'No images found: {args.images}')
    img_paths = img_paths[: args.limit]

    # Load intrinsics
    K, D = load_camera_calib(args.camera_calib)

    # Process
    results: List[Tuple[str, np.ndarray]] = []
    for p in img_paths:
        try:
            q, t = pnp_quaternion_for_image(
                p, K, D,
                args.pattern, args.board_rows, args.board_cols, args.square_size,
                args.charuco_marker,
            )
            results.append((p, q))
            print(f'{os.path.basename(p)}: qx={q[0]:.8f}, qy={q[1]:.8f}, qz={q[2]:.8f}, qw={q[3]:.8f}')
        except Exception as e:
            print(f'WARN: {p}: {e}')

    if args.out:
        import csv
        with open(args.out, 'w', newline='') as f:
            writer = csv.writer(f)
            writer.writerow(['image', 'qx', 'qy', 'qz', 'qw'])
            for p, q in results:
                writer.writerow([os.path.basename(p), f'{q[0]:.8f}', f'{q[1]:.8f}', f'{q[2]:.8f}', f'{q[3]:.8f}'])
        print(f'Written CSV: {args.out}')


if __name__ == '__main__':
    main()
