#!/usr/bin/env python3
"""
readCamPositionTwo.py

Two-camera ArUco detection with joint optimization of both camera extrinsics
against known machine-frame reference markers, plus triangulation of unknown
marker positions. Outputs camera pose and marker poses in machine coordinates,
with CSV and JSON similar to the single-camera script.

Dependencies: numpy, opencv-python (cv2)
Optional but NOT required: SciPy (we implement a simple Levenberg–Marquardt).

Usage example:
  python3 readTwoImages.py -i snapshot_video0_1764531874081.jpg -i snapshot_video1_1764531874081.jpg -npz callibration_cam0.npz -npz callibration_cam1.npz -settings settings.json 
  python3 readTwoImages.py -i snapshot_video0_1764524369655.jpg -i snapshot_video1_1764524369655.jpg -npz callibration_cam0.npz -npz callibration_cam1.npz -settings settings.json 

  python3 readTwoImages.py -i snapshot_video0_1765009029764.jpg -i snapshot_video1_1765009029764.jpg -npz callibration_cam0.npz -npz callibration_cam1.npz -settings settings.json

Settings JSON is expected to contain:
{
  "coordinateSystem": {
    "KnownMarkers": {
       "50": [0.0, 0.0, 0.0],
       "71": [0.140, 0.0, 0.0],
       "101": [0.0, -0.080, 0.0]
    }
  }
}

Author: M365 Copilot (generated)
"""
import argparse
import os
import sys
import json
import time
from typing import Dict, Tuple, List
import numpy as np
import cv2

# ---------------- Configuration defaults ----------------
AXIS_LENGTH_M = 0.05

# ---------------- Utilities ----------------
def load_intrinsics_npz(npz_path: str) -> Tuple[np.ndarray, np.ndarray]:
    print("NPZ reading of file:", npz_path)
    if os.path.exists(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 in NPZ.")
        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)
        print("NPZ loaded:", npz_path)
        return camera_matrix, dist
    # Fallback default intrinsics
    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 get_aruco_detector(dict_name: str):
    mapping = {
        'DICT_4X4_250': cv2.aruco.DICT_4X4_250,
        'DICT_5X5_100': cv2.aruco.DICT_5X5_100,
        'DICT_6X6_250': cv2.aruco.DICT_6X6_250,
        'DICT_ARUCO_ORIGINAL': cv2.aruco.DICT_ARUCO_ORIGINAL,
    }
    if dict_name not in mapping:
        dict_id = cv2.aruco.DICT_4X4_250
    else:
        dict_id = mapping[dict_name]
    dictionary = cv2.aruco.getPredefinedDictionary(dict_id)
    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
    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 corners_to_image_points(corners: np.ndarray) -> np.ndarray:
    return corners.reshape(4,2).astype(np.float32)


def marker_center_from_corners(corners: np.ndarray) -> np.ndarray:
    pts = corners.reshape(4,2)
    return pts.mean(axis=0).astype(np.float32)


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


def rigid_transform_no_scale(A: np.ndarray, B: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
    """Find R,t s.t. B ≈ R A + t. A,B: Nx3."""
    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 undistort_to_normalized(points_px: np.ndarray, K: np.ndarray, D: np.ndarray) -> np.ndarray:
    """points_px: Nx2 pixel. Returns Nx2 normalized coords (x,y) where projection is x=Xp/Z, y=Yp/Z.
    cv2.undistortPoints with P=None yields normalized coordinates.
    """
    pts = points_px.reshape(-1,1,2).astype(np.float32)
    und = cv2.undistortPoints(pts, K, D, P=None)  # returns Nx1x2
    return und.reshape(-1,2)


# ---------------- Joint optimization (LM, numerical Jacobian) ----------------

def pack_params(omega1, t1, omega2, t2) -> np.ndarray:
    return np.hstack([omega1, t1, omega2, t2]).astype(np.float64)


def unpack_params(theta: np.ndarray):
    omega1 = theta[0:3]
    t1     = theta[3:6]
    omega2 = theta[6:9]
    t2     = theta[9:12]
    return omega1, t1, omega2, t2


def residuals_centers_normalized(theta: np.ndarray,
                                 X_mach: np.ndarray,
                                 obs1_norm: np.ndarray,
                                 obs2_norm: np.ndarray) -> np.ndarray:
    """
    Compute residuals in normalized coordinates (no intrinsics, no distortion).
    For camera j: X_cam = R_j X_mach + t_j; proj: (x/z, y/z).
    Returns stacked residuals [r1; r2] shape (4N,), where N = number of references.
    """
    omega1, t1, omega2, t2 = unpack_params(theta)
    R1 = cv2.Rodrigues(omega1)[0]
    R2 = cv2.Rodrigues(omega2)[0]
    # Camera 1 projections
    X_cam1 = (R1 @ X_mach.T + t1.reshape(3,1)).T  # Nx3
    uv1 = X_cam1[:, :2] / X_cam1[:, 2:3]
    r1 = (obs1_norm - uv1).reshape(-1)
    # Camera 2 projections
    X_cam2 = (R2 @ X_mach.T + t2.reshape(3,1)).T
    uv2 = X_cam2[:, :2] / X_cam2[:, 2:3]
    r2 = (obs2_norm - uv2).reshape(-1)
    return np.hstack([r1, r2])


def numerical_jacobian(f, theta: np.ndarray, eps: float, *args) -> np.ndarray:
    """Finite-difference Jacobian: forward differences."""
    r0 = f(theta, *args)
    m = r0.size
    n = theta.size
    J = np.zeros((m, n), dtype=np.float64)
    for k in range(n):
        th = theta.copy()
        th[k] += eps
        rk = f(th, *args)
        J[:, k] = (rk - r0) / eps
    return J, r0


def lm_solve(theta0: np.ndarray,
             X_mach: np.ndarray,
             obs1_norm: np.ndarray,
             obs2_norm: np.ndarray,
             max_iter: int = 50,
             eps_jac: float = 1e-6,
             lambda_init: float = 1e-3) -> Tuple[np.ndarray, Dict]:
    """Simple Levenberg–Marquardt on center normalized residuals."""
    lam = lambda_init
    theta = theta0.copy()
    history = {"iters": [], "rms": []}
    for it in range(max_iter):
        J, r = numerical_jacobian(residuals_centers_normalized, theta, eps_jac,
                                  X_mach, obs1_norm, obs2_norm)
        rms = float(np.sqrt(np.mean(r*r))) if r.size else 0.0
        history["iters"].append(it)
        history["rms"].append(rms)
        # Normal equations with damping
        JTJ = J.T @ J
        g = J.T @ r
        H = JTJ + lam * np.eye(JTJ.shape[0])
        try:
            delta = -np.linalg.solve(H, g)
        except np.linalg.LinAlgError:
            # Fallback to least squares
            delta, *_ = np.linalg.lstsq(H, -g, rcond=None)
        theta_trial = theta + delta
        r_trial = residuals_centers_normalized(theta_trial, X_mach, obs1_norm, obs2_norm)
        rms_trial = float(np.sqrt(np.mean(r_trial*r_trial)))
        if rms_trial < rms:  # improve
            theta = theta_trial
            lam *= 0.5
        else:
            lam *= 2.0
        # Convergence criteria
        if np.linalg.norm(delta) < 1e-9 or abs(rms - rms_trial) < 1e-9:
            break
    return theta, history


# ---------------- Triangulation ----------------

def build_projection_matrix(K: np.ndarray, R: np.ndarray, t: np.ndarray) -> np.ndarray:
    return K @ np.hstack([R, t.reshape(3,1)])


def triangulate_center(P1: np.ndarray, P2: np.ndarray, u1: np.ndarray, u2: np.ndarray) -> np.ndarray:
    # u1,u2: (2,) pixel coordinates
    x1 = u1.reshape(2,1)
    x2 = u2.reshape(2,1)
    X_h = cv2.triangulatePoints(P1, P2, x1, x2)  # 4x1 homogeneous in machine frame if P maps machine->pixels
    X = (X_h[:3] / X_h[3]).reshape(3)
    return X.astype(np.float32)


# ---------------- Main ----------------
def main():
    print("Started")
    parser = argparse.ArgumentParser(description="Two-camera ArUco joint optimization and triangulation")
    parser.add_argument('-i', '--images', action='append', required=True,
                        help="Two image paths: first camera then second camera")
    parser.add_argument('-npz', '--npz', action='append', required=True,
                        help="Two NPZ intrinsics paths: cam1 then cam2")
    parser.add_argument('-markerSizeMM', '--markerSizeMM', type=float, default=25.0,
                        help="Marker side length in millimeters")
    parser.add_argument('--dict', default='DICT_4X4_250', help="ArUco dictionary name")
    parser.add_argument('-settings', type=str, default=None,
                        help="Json settings file containing machine KnownMarkers")
    args = parser.parse_args()

    if len(args.images) != 2 or len(args.npz) != 2:
        print("[ERROR] Provide exactly two images and two intrinsics NPZ files.")
        sys.exit(1)

    img1 = cv2.imread(args.images[0])
    img2 = cv2.imread(args.images[1])
    draw1 = img1.copy()
    draw2 = img2.copy()
    h, w = draw1.shape[:2]
    #drawPNG1 = np.zeros((h, w, 4), dtype=np.uint8)  # fully transparent
    drawPNG1 = np.zeros((h, w, 3), dtype=np.uint8)
    # Also prepare a matching canvas for camera2 to keep the layout tidy
    drawPNG2 = np.zeros((h, w, 3), dtype=np.uint8)

    if img1 is None or img2 is None:
        print("[ERROR] Cannot read one of the images.")
        sys.exit(1)

    K1, D1 = load_intrinsics_npz(args.npz[0])
    K2, D2 = load_intrinsics_npz(args.npz[1])

    # Marker 3D local geometry (square corners)
    half = (args.markerSizeMM / 1000.0) / 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)

    # Read settings for machine known markers
    known_markers: Dict[int, np.ndarray] = {}
    if args.settings is not None and os.path.exists(args.settings):
        with open(args.settings, 'r') as f:
            settings = json.load(f)
            for marker_id, coords in settings['coordinateSystem']['KnownMarkers'].items():
                known_markers[int(marker_id)] = np.array(coords, dtype=np.float32)
        print("[INFO] Loaded KnownMarkers from settings.")
    else:
        # Fallback defaults
        known_markers = {
            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),
        }
        print("[WARN] Using default KnownMarkers; provide -settings for your machine.")

    # Detect markers in both images
    detector_tuple = get_aruco_detector(args.dict)
    corners1, ids1, _ = detect_markers(img1, detector_tuple)
    corners2, ids2, _ = detect_markers(img2, detector_tuple)
    if ids1 is None or ids2 is None:
        print("[ERROR] No markers detected in one or both images.")
        sys.exit(1)

    ids1 = ids1.flatten().tolist()
    ids2 = ids2.flatten().tolist()

    

    # Build dicts: id -> corners, center, rvec/tvec (per-camera PnP)
    def build_marker_dict(img, corners_list, ids, K, D, draw = False) -> Tuple[Dict[int,np.ndarray], Dict[int,np.ndarray], Dict[int,Tuple[np.ndarray,np.ndarray]]]:
        centers = {}
        corners_dict = {}
        poses = {}
        for i, mid in enumerate(ids):
            C = corners_list[i]
            corners_dict[int(mid)] = C
            centers[int(mid)] = marker_center_from_corners(C)
            # Pose from single marker
            img_pts = corners_to_image_points(C)
            success, rvec, tvec = cv2.solvePnP(obj_points, img_pts, K, D, flags=cv2.SOLVEPNP_IPPE_SQUARE)
            if not success:
                success, rvec, tvec = cv2.solvePnP(obj_points, img_pts, K, D)
            if success:
                poses[int(mid)] = (rvec.reshape(3,1), tvec.reshape(3,1))
                if(draw):
                    cv2.drawFrameAxes(draw1, K, D, rvec, tvec, 0.02)  # slim orientation axes
                    cv2.drawFrameAxes(drawPNG1, K, D, rvec, tvec, 0.02)  # slim orientation axes
        return centers, corners_dict, poses

    centers1, corners_dict1, poses1 = build_marker_dict(img1, corners1, ids1, K1, D1, draw = True)
    centers2, corners_dict2, poses2 = build_marker_dict(img2, corners2, ids2, K2, D2)

    # Common reference markers present in both images and known in machine frame
    common_refs = [mid for mid in known_markers.keys() if (mid in centers1 and mid in centers2)]
    if len(common_refs) < 3:
        print(f"[ERROR] Need ≥3 common reference markers in both cameras; found {len(common_refs)}: {common_refs}")
        sys.exit(1)

    # Initial extrinsics from rigid fits per camera using tvec centers of references
    # For camera j, A = camera-frame positions of reference markers (from PnP tvec), B = machine positions
    def init_extrinsics_from_rigid(poses_cam: Dict[int,Tuple[np.ndarray,np.ndarray]], refs: List[int]) -> Tuple[np.ndarray,np.ndarray]:
        A = []
        B = []
        for mid in refs:
            if mid in poses_cam:
                _, tvec = poses_cam[mid]
                A.append(tvec.flatten())
                B.append(known_markers[mid])
        if len(A) < 2:
            raise RuntimeError("Insufficient reference poses for rigid transform init.")
        A = np.stack(A, axis=0)
        B = np.stack(B, axis=0)
        R_cm, t_cm = rigid_transform_no_scale(A, B)  # camera->machine
        # Convert to machine->camera
        R_mc = R_cm.T
        t_mc = - R_mc @ t_cm
        return R_mc.astype(np.float32), t_mc.astype(np.float32)

    R1_init, t1_init = init_extrinsics_from_rigid(poses1, common_refs)
    R2_init, t2_init = init_extrinsics_from_rigid(poses2, common_refs)

    # Observations: reference centers (pixels) -> normalized
    X_mach_refs = np.stack([known_markers[mid] for mid in common_refs], axis=0).astype(np.float32)
    obs1_px = np.stack([centers1[mid] for mid in common_refs], axis=0).astype(np.float32)
    obs2_px = np.stack([centers2[mid] for mid in common_refs], axis=0).astype(np.float32)
    obs1_norm = undistort_to_normalized(obs1_px, K1, D1)
    obs2_norm = undistort_to_normalized(obs2_px, K2, D2)

    # Pack initial params as Rodrigues vectors
    omega1_init = cv2.Rodrigues(R1_init)[0].reshape(3)
    omega2_init = cv2.Rodrigues(R2_init)[0].reshape(3)
    theta0 = pack_params(omega1_init, t1_init.reshape(3), omega2_init, t2_init.reshape(3))

    theta_opt, hist = lm_solve(theta0, X_mach_refs, obs1_norm, obs2_norm,
                               max_iter=60, eps_jac=1e-6, lambda_init=1e-3)

    omega1, t1, omega2, t2 = unpack_params(theta_opt)
    R1_opt = cv2.Rodrigues(omega1)[0]
    R2_opt = cv2.Rodrigues(omega2)[0]

    # Camera pose in machine coordinates (camera->machine): R_cm = R^T, t_cm = -R^T t
    R1_cm = R1_opt.T
    t1_cm = - R1_cm @ t1
    R2_cm = R2_opt.T
    t2_cm = - R2_cm @ t2

    # Build projection matrices for triangulation (machine -> pixels)
    P1 = build_projection_matrix(K1, R1_opt, t1)
    P2 = build_projection_matrix(K2, R2_opt, t2)

    # Collect markers seen by at least one camera
    all_ids = set(ids1) | set(ids2)
    # Output structures
    rows = [("id", "x_mm", "y_mm", "z_mm", "roll_deg", "pitch_deg", "yaw_deg")]
    marker_list = []

    # Camera orientations in Euler (ZYX)
    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

    cam1_roll, cam1_pitch, cam1_yaw = R_to_euler_zyx(R1_cm)
    cam2_roll, cam2_pitch, cam2_yaw = R_to_euler_zyx(R2_cm)

    # Camera rows
    c1_mm = (t1_cm * 1000.0).tolist()
    rows.append(("camera 0", f"{c1_mm[0]:.2f}", f"{c1_mm[1]:.2f}", f"{c1_mm[2]:.2f}", f"{cam1_roll:.3f}", f"{cam1_pitch:.3f}", f"{cam1_yaw:.3f}"))
    c2_mm = (t2_cm * 1000.0).tolist()
    rows.append(("camera 1", f"{c2_mm[0]:.2f}", f"{c2_mm[1]:.2f}", f"{c2_mm[2]:.2f}", f"{cam2_roll:.3f}", f"{cam2_pitch:.3f}", f"{cam2_yaw:.3f}"))

    # Triangulate and output markers
    def orientation_in_machine(mid: int) -> Tuple[float,float,float]:
        # Prefer camera1 orientation, else camera2
        if mid in poses1:
            Rm_cam = rvec_to_R(poses1[mid][0])
            Rm_machine = R1_cm @ Rm_cam
        elif mid in poses2:
            Rm_cam = rvec_to_R(poses2[mid][0])
            Rm_machine = R2_cm @ Rm_cam
        else:
            Rm_machine = np.eye(3, dtype=np.float32)
        return R_to_euler_zyx(Rm_machine)

    # Residual report for references
    # Recompute residual RMS in pixels for references (after optimization)
    refs_rms_px = []
    for j,(K,R_opt,t_opt,centers_px) in enumerate([(K1,R1_opt,t1,centers1),(K2,R2_opt,t2,centers2)]):
        errs = []
        for mid in common_refs:
            X = known_markers[mid]
            X_cam = R_opt @ X + t_opt
            x = K @ X_cam
            u_pred = x[0]/x[2]
            v_pred = x[1]/x[2]
            u_obs, v_obs = centers_px[mid]
            errs.append(np.hypot(u_obs-u_pred, v_obs-v_pred))
        refs_rms_px.append(float(np.sqrt(np.mean(np.square(errs))) if errs else 0.0))

    # Compute per-marker positions
    for mid in sorted(all_ids):
        # Triangulate if seen in both
        if mid in centers1 and mid in centers2:
            X_mach = triangulate_center(P1, P2, centers1[mid], centers2[mid])
        elif mid in poses1:
            # Fallback single-camera: transform tvec (camera->machine)
            X_mach = (R1_cm @ poses1[mid][1].flatten() + t1_cm)
        elif mid in poses2:
            X_mach = (R2_cm @ poses2[mid][1].flatten() + t2_cm)
        else:
            continue
        roll, pitch, yaw = orientation_in_machine(mid)
        x_mm, y_mm, z_mm = (X_mach * 1000.0).tolist()
        rows.append((mid, f"{x_mm:.2f}", f"{y_mm:.2f}", f"{z_mm:.2f}", f"{roll:.3f}", f"{pitch:.3f}", f"{yaw:.3f}"))
        marker_list.append({
            "id": int(mid),
            "position_mm": [float(x_mm), float(y_mm), float(z_mm)],
            "orientation_deg": {"roll": float(roll), "pitch": float(pitch), "yaw": float(yaw)}
        })

    # Save CSV & JSON
    base1 = os.path.splitext(args.images[0])[0]
    base2 = os.path.splitext(args.images[1])[0]
    out_csv = f"{base1}_two_cam.csv"
    out_json = f"{base1}_two_cam.json"

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

    json_data = {
        "metadata": {
            "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
            "reference_markers": common_refs,
            "dict": args.dict,
            "marker_size_mm": args.markerSizeMM,
            "rms_refs_px_cam1": refs_rms_px[0] if refs_rms_px else None,
            "rms_refs_px_cam2": refs_rms_px[1] if refs_rms_px else None,
            "description": "Two-camera joint optimization with triangulation"
        },
        "cameras": [
            {
                "id": "camera1",
                "position_mm": [float(x) for x in (t1_cm * 1000.0)],
                "orientation_deg": {"roll": cam1_roll, "pitch": cam1_pitch, "yaw": cam1_yaw},
            },
            {
                "id": "camera2",
                "position_mm": [float(x) for x in (t2_cm * 1000.0)],
                "orientation_deg": {"roll": cam2_roll, "pitch": cam2_pitch, "yaw": cam2_yaw},
            }
        ],
        "markers": marker_list
    }

    try:
        with open(out_json, 'w', encoding='utf-8') as f:
            json.dump(json_data, f, indent=4)
        print(f"[INFO] Saved JSON poses to '{out_json}'.")
    except Exception as e:
        print(f"[WARN] Could not save JSON: {e}")

    # Annotate images with machine axes using camera1 extrinsics
    try:
        R_machine_to_cam1 = R1_opt
        t_machine_to_cam1 = t1
        machine_axes = np.float32([
            [0.0, 0.0, 0.0],
            [0.200, 0.0, 0.0],
            [0.0, -0.100, 0.0],
            [0.0, 0.0, 0.100],
        ])
        rvec_global, _ = cv2.Rodrigues(R_machine_to_cam1)
        imgpts, _ = cv2.projectPoints(machine_axes, rvec_global, t_machine_to_cam1, K1, D1)
        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))
        
        # Draw marker outlines only (omit default small id labels) — we draw larger IDs below
        cv2.aruco.drawDetectedMarkers(draw1, corners1)
        cv2.aruco.drawDetectedMarkers(drawPNG1, corners1)
        # Draw larger blue ID labels (keep default marker outlines as-is)
        try:
            for i, mid in enumerate(ids1):
                try:
                    pts = corners1[i].reshape((4, 2))
                    center = tuple(np.round(pts.mean(axis=0)).astype(int))
                except Exception:
                    continue
                text = str(int(mid))
                # Offset: 5px more to the right and 5px up (y axis is downwards)
                pos = (int(center[0]) + 15, int(center[1]) - 15)
                cv2.putText(draw1, text, pos, cv2.FONT_HERSHEY_SIMPLEX, 1.0, (255,0,0), 3, lineType=cv2.LINE_AA)
                cv2.putText(drawPNG1, text, pos, cv2.FONT_HERSHEY_SIMPLEX, 1.0, (255,0,0,255), 3, lineType=cv2.LINE_AA)
        except Exception:
            pass
        cv2.line(draw1, origin, x_end, (0,0,255), 3)
        cv2.line(draw1, origin, y_end, (0,255,0), 3)
        cv2.line(draw1, origin, z_end, (255,0,0), 3)
        
        # Draw lines (RGBA colors: B,G,R,A). A=255 for fully opaque.
        cv2.line(drawPNG1, origin, x_end, (0,   0, 255, 255), 3)  # Red X
        cv2.line(drawPNG1, origin, y_end, (0, 255,   0, 255), 3)  # Green Y
        cv2.line(drawPNG1, origin, z_end, (255, 0,   0, 255), 3)  # Blue Z

        cv2.putText(draw1, "X (200 mm)", x_end, cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0,0,255), 2)
        cv2.putText(draw1, "Y (-100 mm)", y_end, cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0,255,0), 2)
        cv2.putText(draw1, "+Z (100 mm)", z_end, cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255,0,0), 2)
        
        
        # Try to draw text (might be jaggy on transparent BG in older OpenCV)
        cv2.putText(drawPNG1, "X (200 mm)", x_end, cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0,   0, 255, 255), 2)
        cv2.putText(drawPNG1, "Y (-100 mm)", y_end, cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255,   0, 255), 2)
        cv2.putText(drawPNG1, "+Z (100 mm)", z_end, cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 0,   0, 255), 2)


        out_img1 = f"{base1}_two_cam_annotated.jpg"
        cv2.imwrite(out_img1, draw1)
        print(f"[INFO] Annotated image saved as '{out_img1}'.")
        
        # Save as transparent PNG
        
        gray = cv2.cvtColor(drawPNG1, cv2.COLOR_BGR2GRAY)
        _, alpha = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY)

        # 5) Merge BGR + alpha → RGBA transparent overlay
        drawPNG_1 = cv2.merge([drawPNG1[:, :, 0], drawPNG1[:, :, 1], drawPNG1[:, :, 2], alpha])

        out_png1 = f"{base1}_two_cam_overlay.png"
        cv2.imwrite(out_png1, drawPNG_1)

    except Exception as e:
        print(f"[WARN] Failed to draw machine axes: {e}")

    # Also annotate the second camera image and produce a transparent overlay
    try:
        machine_axes2 = np.float32([
            [0.0, 0.0, 0.0],
            [0.200, 0.0, 0.0],
            [0.0, -0.100, 0.0],
            [0.0, 0.0, 0.100],
        ])
        rvec_global2, _ = cv2.Rodrigues(R2_opt)
        imgpts2, _ = cv2.projectPoints(machine_axes2, rvec_global2, t2, K2, D2)
        origin2 = tuple(np.round(imgpts2[0].ravel()).astype(int))
        x_end2 = tuple(np.round(imgpts2[1].ravel()).astype(int))
        y_end2 = tuple(np.round(imgpts2[2].ravel()).astype(int))
        z_end2 = tuple(np.round(imgpts2[3].ravel()).astype(int))

        # Draw marker outlines only (omit default small id labels) — we draw larger IDs below
        cv2.aruco.drawDetectedMarkers(draw2, corners2)
        cv2.aruco.drawDetectedMarkers(drawPNG2, corners2)
        # Draw larger blue ID labels (keep default marker outlines as-is)
        try:
            for i, mid in enumerate(ids2):
                try:
                    pts = corners2[i].reshape((4, 2))
                    center = tuple(np.round(pts.mean(axis=0)).astype(int))
                except Exception:
                    continue
                text = str(int(mid))
                # Offset: 5px more to the right and 5px up (y axis is downwards)
                pos = (int(center[0]) + 13, int(center[1]) + 3)
                cv2.putText(draw2, text, pos, cv2.FONT_HERSHEY_SIMPLEX, 1.0, (255,0,0), 3, lineType=cv2.LINE_AA)
                cv2.putText(drawPNG2, text, pos, cv2.FONT_HERSHEY_SIMPLEX, 1.0, (255,0,0,255), 3, lineType=cv2.LINE_AA)
        except Exception:
            pass

        cv2.line(draw2, origin2, x_end2, (0,0,255), 3)
        cv2.line(draw2, origin2, y_end2, (0,255,0), 3)
        cv2.line(draw2, origin2, z_end2, (255,0,0), 3)

        cv2.line(drawPNG2, origin2, x_end2, (0,   0, 255, 255), 3)
        cv2.line(drawPNG2, origin2, y_end2, (0, 255,   0, 255), 3)
        cv2.line(drawPNG2, origin2, z_end2, (255, 0,   0, 255), 3)

        cv2.putText(draw2, "X (200 mm)", x_end2, cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0,0,255), 2)
        cv2.putText(draw2, "Y (-100 mm)", y_end2, cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0,255,0), 2)
        cv2.putText(draw2, "+Z (100 mm)", z_end2, cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255,0,0), 2)

        cv2.putText(drawPNG2, "X (200 mm)", x_end2, cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0,   0, 255, 255), 2)
        cv2.putText(drawPNG2, "Y (-100 mm)", y_end2, cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255,   0, 255), 2)
        cv2.putText(drawPNG2, "+Z (100 mm)", z_end2, cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 0,   0, 255), 2)

        out_img2 = f"{base2}_two_cam_annotated.jpg"
        cv2.imwrite(out_img2, draw2)
        print(f"[INFO] Annotated image saved as '{out_img2}'.")

        gray2 = cv2.cvtColor(drawPNG2, cv2.COLOR_BGR2GRAY)
        _, alpha2 = cv2.threshold(gray2, 0, 255, cv2.THRESH_BINARY)
        drawPNG_2 = cv2.merge([drawPNG2[:, :, 0], drawPNG2[:, :, 1], drawPNG2[:, :, 2], alpha2])
        out_png2 = f"{base2}_two_cam_overlay.png"
        cv2.imwrite(out_png2, drawPNG_2)
        print(f"[INFO] Overlay PNG saved as '{out_png2}'.")

    except Exception as e:
        print(f"[WARN] Failed to draw machine axes for camera2: {e}")

if __name__ == '__main__':
    main()
