#!/usr/bin/env python3
"""
ArUco detection with multi-marker machine-frame fit + camera pose output (OpenCV >= 4.8).

- Reads: webCam_1.jpg
- Detects DICT_4X4_250 markers (ids expected: 0, 5, 10, 15)
- Uses multiple reference markers with known machine coordinates to fit camera->machine transform
- Reports positions/orientations of all markers **and the camera** in machine coordinates
- Draws detected markers, per-marker axes, and the machine axes
- Saves: webCam_1a.jpg (annotated) and marker_poses_machine.csv (poses incl. camera)


Usage:
   python3 readCamPos.py -i snapshot_video1_1764493534200.jpg -npz camera_intrinsics_v0.npz -setting settings.json

   
"""
import faulthandler
faulthandler.enable()

import argparse
import os
import sys
import csv
import json
import time
from typing import Tuple, Dict, List
import numpy as np
import cv2



# ----------------------- Configuration Defaults -----------------------
IMAGE_PATH = "default.jpg"
OUTPUT_IMAGE_PATH = "default.jpg"
OUTPUT_CSV_PATH = "default.csv"
OUTPUT_JSON_PATH = "default.json"

# Marker side length in meters (25 mm)
MARKER_LENGTH_M = 0.025

# Axis lengths for visualization (in meters)
AXIS_LENGTH_M = 0.05           # per-marker axis
MACHINE_AXIS_X_M = 0.200       # 200 mm along +X
MACHINE_AXIS_Y_M = -0.100      # -100 mm along Y (towards camera per description)
MACHINE_AXIS_Z_M = 0.100       # +Z visualized as 100 mm

# Known machine coordinates for reference markers (meters)
cam_anchor_pts = {}

EXPECTED_IDS = {50, 71, 101}

# ----------------------- Utilities -----------------------

def load_intrinsics_npz(npz_path: str) -> Tuple[np.ndarray, np.ndarray]:
    if os.path.exists(npz_path):
        
        print("NPZ from File:", npz_path)
        data = np.load(npz_path)
        for k in ('camera_matrix', 'mtx', 'K'):
            if k in data:
                camera_matrix = data[k].astype(np.float32)
                break
        else:
            raise KeyError("Camera matrix not found.")
        for k in ('dist_coeffs', 'dist', 'D'):
            if k in data:
                dist = data[k].astype(np.float32).reshape(-1, 1)
                break
        else:
            dist = np.zeros((5, 1), dtype=np.float32)
        return camera_matrix, dist

    camera_matrix = np.array([[1400, 0, 640],
                              [0, 1400, 360],
                              [0, 0, 1]], dtype=np.float32)
    
    dist_coeffs = np.zeros((5, 1), dtype=np.float32)
    
    print("[WARN] Using default approximate intrinsics.")
    return camera_matrix, dist_coeffs


def rvec_to_R(rvec: np.ndarray) -> np.ndarray:
    R, _ = cv2.Rodrigues(rvec)
    return R


def R_to_euler_zyx(R: np.ndarray) -> Tuple[float, float, float]:
    yaw = float(np.degrees(np.arctan2(R[1,0], R[0,0])))
    sp = np.sqrt(R[2,1]**2 + R[2,2]**2)
    pitch = float(np.degrees(np.arctan2(-R[2,0], sp)))
    roll = float(np.degrees(np.arctan2(R[2,1], R[2,2])))
    return roll, pitch, yaw


def corners_to_image_points(corners: np.ndarray) -> np.ndarray:
    return corners.reshape(4, 2).astype(np.float32)


def get_detector():
    dictionary = cv2.aruco.getPredefinedDictionary(cv2.aruco.DICT_4X4_250)
    try:
        params = cv2.aruco.DetectorParameters()
    except Exception:
        params = cv2.aruco.DetectorParameters_create()
    try:
        detector = cv2.aruco.ArucoDetector(dictionary, params)
        return detector, None
    except Exception:
        return None, (dictionary, params)


def detect_markers(image: np.ndarray, detector_tuple):
    detector, fallback = detector_tuple
    print(detector)
    if detector is not None:
        corners, ids, rejected = detector.detectMarkers(image)
    else:
        dictionary, params = fallback
        corners, ids, rejected = cv2.aruco.detectMarkers(image, dictionary, parameters=params)
    return corners, ids, rejected


def rigid_transform_no_scale(A: np.ndarray, B: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
    assert A.shape == B.shape and A.shape[1] == 3, "A and B must be Nx3"
    N = A.shape[0]
    if N < 2:
        raise ValueError("Need at least 2 points; 3+ recommended.")
    centroid_A = A.mean(axis=0)
    centroid_B = B.mean(axis=0)
    AA = A - centroid_A
    BB = B - centroid_B
    H = AA.T @ BB
    U, S, Vt = np.linalg.svd(H)
    R = Vt.T @ U.T
    if np.linalg.det(R) < 0:
        Vt[-1, :] *= -1
        R = Vt.T @ U.T
    t = centroid_B - R @ centroid_A
    return R.astype(np.float32), t.astype(np.float32)

def readSettings(fileSetting):
    global cam_anchor_pts
    print("Read Settings")

    if(fileSetting == None):
        cam_anchor_pts = {
            50: np.array([0.0,   0.0,   0.0], dtype=np.float32),
            71: np.array([0.140, 0.0,   0.0], dtype=np.float32),
            101: np.array([-0.0, -0.080, 0.0], dtype=np.float32),
            #15: np.array([20,20,20])  # add if known
        }
        return
    

    with open(fileSetting, 'r') as f:
        settings = json.load(f)
    for marker_id, coords in settings['coordinateSystem']['KnownMarkers'].items():
        cam_anchor_pts[int(marker_id)] = np.array(coords, dtype=np.float32)

    #KNOWN_MACHINE_POS = {int(k): np.array(v, dtype=np.float32) for k, v in settings.items()}
    

# ----------------------- Main -----------------------

def main():

    parser = argparse.ArgumentParser(description="Detect ArUco markers in two images and compute camera poses in machine coordinates.")
    parser.add_argument('-i', '--images', action='append', required=False,
                       help="Path to image. Provide this option twice: once per camera (e.g., -i2 cam1.jpg -i2 cam2.jpg)")
    parser.add_argument('-npz', '--npz', action='append', required=False, default=['camera_intrinsics_v1.npz'])
    parser.add_argument('--cam-calib', action='append', required=False,
        help="Paths to calibration YAMLs for camera 1 and camera 2 (e.g., cam1.npz cam2.npz)")
    parser.add_argument('--marker-size-mm', type=float, default=25,
        help="Marker side length in millimeters (e.g., 50)")
    parser.add_argument('--dict', default='DICT_4X4_250',
        help="ArUco dictionary name (default: DICT_4X4_250)")
    parser.add_argument('-settings', type=str, default=None,
        help="Json File with Machine Settings")
    args = parser.parse_args()
    

    
    print("ABC 0")

    readSettings(args.settings)
    
    print("ABC 0")

    if(args.images is None):
         image = cv2.imread(IMAGE_PATH)
    else:
        image = cv2.imread(args.images[0])
        OUTPUT_IMAGE_PATH = args.images[0].replace(".jpg","r.jpg").replace(".PNG","r.PNG")
        OUTPUT_CSV_PATH = args.images[0].replace(".jpg",".csv").replace(".PNG",".csv")
        OUTPUT_JSON_PATH = args.images[0].replace(".jpg",".json").replace(".PNG",".json")


    if image is None:
        print(f"[ERROR] Cannot read image '{IMAGE_PATH}'.")
        sys.exit(1)


    print("ABC 1")

    camera_matrix, dist_coeffs = load_intrinsics_npz(args.npz[0])
    print("ABC 1a")
    detector_tuple = get_detector()
    print("ABC 1b")
    corners_list, ids, rejected = detect_markers(image, detector_tuple)


    print("ABC 2")

    if ids is None or len(ids) == 0:
        print("[ERROR] No markers detected.")
        sys.exit(1)

    draw_img = image.copy()
    cv2.aruco.drawDetectedMarkers(draw_img, corners_list, ids)

    half = MARKER_LENGTH_M / 2.0
    obj_points = np.array([
        [-half,  half, 0.0],
        [ half,  half, 0.0],
        [ half, -half, 0.0],
        [-half, -half, 0.0],
    ], dtype=np.float32)

    poses_cam: Dict[int, Tuple[np.ndarray, np.ndarray]] = {}
    centers_cam: Dict[int, np.ndarray] = {}
    for i, marker_id in enumerate(ids.flatten()):
        img_pts = corners_to_image_points(corners_list[i])
        success, rvec, tvec = cv2.solvePnP(obj_points, img_pts, camera_matrix, dist_coeffs, flags=cv2.SOLVEPNP_IPPE_SQUARE)
        if not success:
            success, rvec, tvec = cv2.solvePnP(obj_points, img_pts, camera_matrix, dist_coeffs)
        if success:
            rvec = rvec.reshape(3,1)
            tvec = tvec.reshape(3,1)
            poses_cam[int(marker_id)] = (rvec, tvec)
            centers_cam[int(marker_id)] = tvec.flatten()
            try:
                cv2.drawFrameAxes(draw_img, camera_matrix, dist_coeffs, rvec, tvec, AXIS_LENGTH_M)
            except Exception:
                pass
        else:
            print(f"[WARN] solvePnP failed for marker {marker_id}")

    common_ids: List[int] = [mid for mid in cam_anchor_pts.keys() if mid in centers_cam]
    if len(common_ids) < 2:
        print(f"[ERROR] Need at least 2 reference markers; found {len(common_ids)}: {common_ids}")
        sys.exit(1)
    if len(common_ids) < 3:
        print(f"[WARN] Only {len(common_ids)} references ({common_ids}). Fit may be less stable; 3+ recommended.")

    A = np.stack([centers_cam[mid] for mid in common_ids], axis=0)
    B = np.stack([cam_anchor_pts[mid] for mid in common_ids], axis=0)

    R_cam_to_machine, t_cam_to_machine = rigid_transform_no_scale(A, B)

    residuals_mm = []
    for i, mid in enumerate(common_ids):
        pred = R_cam_to_machine @ A[i] + t_cam_to_machine
        err = np.linalg.norm(pred - B[i]) * 1000.0
        residuals_mm.append(err)
    rms = float(np.sqrt(np.mean(np.square(residuals_mm)))) if residuals_mm else 0.0
    print("\nReference fit residuals (mm) per marker:")
    for mid, e in zip(common_ids, residuals_mm):
        print(f"  ID {mid}: {e:.2f} mm")
    print(f"RMS residual: {rms:.2f} mm")

    # Camera pose in machine coordinates:
    # Camera origin (0,0,0 in camera) maps to t_cam_to_machine
    cam_pos_machine = t_cam_to_machine
    cam_R_machine = R_cam_to_machine  # camera basis expressed in machine frame
    cam_roll, cam_pitch, cam_yaw = R_to_euler_zyx(cam_R_machine)

    rows = [("id", "x_mm", "y_mm", "z_mm", "roll_deg", "pitch_deg", "yaw_deg")]
    marker_list = []

    print("\nMarker Positions and Orientations in Machine Coordinates:")
    print(f"{'ID':>8}  {'X(mm)':>10}  {'Y(mm)':>10}  {'Z(mm)':>10}  {'Roll':>10}  {'Pitch':>10}  {'Yaw':>10}")

    # Add camera first
    cx, cy, cz = (cam_pos_machine * 1000.0).tolist()
    print(f"{'camera':>8}  {cx:10.2f}  {cy:10.2f}  {cz:10.2f}  {cam_roll:10.2f}  {cam_pitch:10.2f}  {cam_yaw:10.2f}")
    rows.append(("camera", f"{cx:.3f}", f"{cy:.3f}", f"{cz:.3f}", f"{cam_roll:.3f}", f"{cam_pitch:.3f}", f"{cam_yaw:.3f}"))
    camera_pose = {
    "id": "camera",
    "position_mm": [float(x) for x in cam_pos_machine * 1000.0],
    "orientation_deg": {"roll": cam_roll, "pitch": cam_pitch, "yaw": cam_yaw}
}

    # Then markers
    for marker_id in sorted(poses_cam.keys()):
        rvec, tvec = poses_cam[marker_id]
        R_marker_cam = rvec_to_R(rvec)
        pos_machine = R_cam_to_machine @ tvec.flatten() + t_cam_to_machine
        R_marker_machine = R_cam_to_machine @ R_marker_cam
        roll_deg, pitch_deg, yaw_deg = R_to_euler_zyx(R_marker_machine)
        x_mm, y_mm, z_mm = (pos_machine * 1000.0).tolist()
        print(f"{marker_id:8d}  {x_mm:10.2f}  {y_mm:10.2f}  {z_mm:10.2f}  {roll_deg:10.2f}  {pitch_deg:10.2f}  {yaw_deg:10.2f}")
        rows.append((marker_id, f"{x_mm:.3f}", f"{y_mm:.3f}", f"{z_mm:.3f}", f"{roll_deg:.3f}", f"{pitch_deg:.3f}", f"{yaw_deg:.3f}"))
        marker_list.append({"id": marker_id, "position_mm": [x_mm, y_mm, z_mm], "orientation_deg": {"roll": roll_deg, "pitch": pitch_deg, "yaw": yaw_deg}})

    # Save CSV
    try:
        with open(OUTPUT_CSV_PATH, 'w', newline='') as f:
            writer = csv.writer(f)
            writer.writerows(rows)
        print(f"\n[INFO] Saved CSV poses to '{OUTPUT_CSV_PATH}'.")
    except Exception as e:
        print(f"[WARN] Could not save CSV: {e}")


    # Save JSON
    json_data = {
        "metadata": {
            "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
            "reference_markers": common_ids,
            "rms_residual_mm": rms,
            "description": "Multi-marker machine frame fit with camera pose"
        },
        "camera": camera_pose,
        "markers": marker_list
    }
    with open(OUTPUT_JSON_PATH, 'w', encoding='utf-8') as f:
        json.dump(json_data, f, indent=4)  

    # Warn about expected IDs
    detected_ids = set(poses_cam.keys())
    missing = EXPECTED_IDS - detected_ids
    if missing:
        print(f"[WARN] Expected markers not detected: {sorted(missing)}")

    # Draw machine axes using global transform (machine->camera)
    R_machine_to_cam = R_cam_to_machine.T
    t_machine_to_cam = - R_machine_to_cam @ t_cam_to_machine
    try:
        machine_axes = np.float32([
            [0.0, 0.0, 0.0],
            [MACHINE_AXIS_X_M, 0.0, 0.0],
            [0.0, MACHINE_AXIS_Y_M, 0.0],
            [0.0, 0.0, MACHINE_AXIS_Z_M],
        ])
        rvec_global, _ = cv2.Rodrigues(R_machine_to_cam)
        imgpts, _ = cv2.projectPoints(machine_axes, rvec_global, t_machine_to_cam, camera_matrix, dist_coeffs)
        origin = tuple(np.round(imgpts[0].ravel()).astype(int))
        x_end = tuple(np.round(imgpts[1].ravel()).astype(int))
        y_end = tuple(np.round(imgpts[2].ravel()).astype(int))
        z_end = tuple(np.round(imgpts[3].ravel()).astype(int))
        cv2.line(draw_img, origin, x_end, (0, 0, 255), 3)
        cv2.line(draw_img, origin, y_end, (0, 255, 0), 3)
        cv2.line(draw_img, origin, z_end, (255, 0, 0), 3)
        cv2.putText(draw_img, "X (200 mm)", x_end, cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0,0,255), 2)
        cv2.putText(draw_img, "Y (-100 mm)", y_end, cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0,255,0), 2)
        cv2.putText(draw_img, "+Z (100 mm)", z_end, cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255,0,0), 2)
    except Exception as e:
        print(f"[WARN] Failed to draw machine axes: {e}")

    ok = cv2.imwrite(OUTPUT_IMAGE_PATH, draw_img)
    if ok:
        print(f"[INFO] Annotated image saved as '{OUTPUT_IMAGE_PATH}'.")
    else:
        print(f"[ERROR] Failed to save annotated image '{OUTPUT_IMAGE_PATH}'.")

if __name__ == '__main__':
    main()
