#!/usr/bin/env python3

import argparse
import json
import numpy as np
import cv2


# ------------------------------------------------------------
# Load JSON
# ------------------------------------------------------------

def load_json(path):
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)


# ------------------------------------------------------------
# Robot model: marker centers in world coordinates
# ------------------------------------------------------------

def load_robot_markers(robot_json):
    markers = {}
    for m in robot_json["Marker"]:
        if m.get("on") == "Base":
            mid = int(m["id"])
            markers[mid] = np.array(m["relPos"], dtype=np.float32)
    return markers


# ------------------------------------------------------------
# Marker geometry (world frame)
# ------------------------------------------------------------

def marker_corners_world(center, size_m):
    """
    Returns 4 corners in consistent OpenCV order:
    TL, TR, BR, BL
    Marker lies in XY plane (z=0)
    """
    h = size_m / 2.0
    x, y, z = center

    return np.array([
        [x - h, y + h, z],
        [x + h, y + h, z],
        [x + h, y - h, z],
        [x - h, y - h, z],
    ], dtype=np.float32)


# ------------------------------------------------------------
# Build correspondences for one camera
# ------------------------------------------------------------

def build_correspondences(camera, robot_markers, marker_size_m):

    obj_pts = []
    img_pts = []

    for obs in camera["observations"]:
        mid = int(obs["marker_id"])

        if mid not in robot_markers:
            continue

        center = robot_markers[mid]

        obj_corners = marker_corners_world(center, marker_size_m)
        img_corners = np.array(obs["corners_px"], dtype=np.float32)

        obj_pts.append(obj_corners)
        img_pts.append(img_corners)

    if len(obj_pts) == 0:
        return None, None

    obj_pts = np.vstack(obj_pts).astype(np.float32)
    img_pts = np.vstack(img_pts).astype(np.float32)

    return obj_pts, img_pts


# ------------------------------------------------------------
# Solve PnP
# ------------------------------------------------------------

def solve_camera(obj_pts, img_pts, K, dist):

    if obj_pts is None or len(obj_pts) < 6:
        raise RuntimeError("Not enough correspondences for PnP")

    ok, rvec, tvec = cv2.solvePnP(
        obj_pts,
        img_pts,
        K,
        dist,
        flags=cv2.SOLVEPNP_ITERATIVE
    )

    if not ok:
        raise RuntimeError("solvePnP failed")

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


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

def main():

    parser = argparse.ArgumentParser()

    parser.add_argument("-scene", required=True)
    parser.add_argument("-robot", required=True)
    parser.add_argument("--marker_size", type=float, default=0.025)
    parser.add_argument("-out", default="camera_poses.json")

    args = parser.parse_args()

    scene = load_json(args.scene)
    robot = load_json(args.robot)

    robot_markers = load_robot_markers(robot)

    result = {
        "camera_poses": {}
    }

    # --------------------------------------------------------
    # Each camera independently
    # --------------------------------------------------------

    for cam_id, cam in scene["cameras"].items():

        print(f"[INFO] Solving camera {cam_id}")

        K = np.array(cam["camera_matrix"], dtype=np.float32)
        dist = np.array(cam["distortion_coefficients"], dtype=np.float32)

        obj_pts, img_pts = build_correspondences(
            cam,
            robot_markers,
            args.marker_size
        )

        if obj_pts is None:
            print(f"[WARN] Camera {cam_id}: no valid markers")
            continue

        R, t = solve_camera(obj_pts, img_pts, K, dist)

        result["camera_poses"][cam_id] = {
            "R_world_from_cam": R.tolist(),
            "t_world_from_cam": t.flatten().tolist()
        }

        print(f"[OK] Camera {cam_id} solved")

    # --------------------------------------------------------
    # Save
    # --------------------------------------------------------

    with open(args.out, "w", encoding="utf-8") as f:
        json.dump(result, f, indent=2)

    print(f"[DONE] Saved -> {args.out}")


if __name__ == "__main__":
    main()