#!/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") != "Board":
            continue
        if "id" not in m:
            continue
        pos = m.get("position")
        if pos is None:
            pos = m.get("relPos")
        if pos is None:
            continue
        mid = int(m["id"])
        markers[mid] = np.array(pos, 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)


def marker_corners_local(size_m):
    h = size_m / 2.0
    return np.array([
        [-h,  h, 0.0],
        [ h,  h, 0.0],
        [ h, -h, 0.0],
        [-h, -h, 0.0],
    ], dtype=np.float32)


# ------------------------------------------------------------
# Solve single marker pose
# ------------------------------------------------------------

def solve_marker_pose(corners_px, K, dist, marker_size_m):
    obj_pts = marker_corners_local(marker_size_m)
    ok, rvec, tvec = cv2.solvePnP(
        obj_pts,
        corners_px,
        K,
        dist,
        flags=cv2.SOLVEPNP_IPPE_SQUARE
    )

    if not ok:
        ok, rvec, tvec = cv2.solvePnP(
            obj_pts,
            corners_px,
            K,
            dist,
            flags=cv2.SOLVEPNP_ITERATIVE
        )

    if not ok:
        return None, None

    return rvec, tvec


def rigid_transform_no_scale(A: np.ndarray, B: np.ndarray):
    """Find R, t such that B ≈ R A + t for A,B: Nx3."""
    assert A.shape == B.shape and A.shape[1] == 3, "A and B must be Nx3"
    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)


# ------------------------------------------------------------
# Estimate camera pose from board marker center correspondences
# ------------------------------------------------------------

def build_camera_pose_from_board_markers(camera_id, scene_markers, robot_markers, K, dist, marker_size_m):
    cam_centers = []
    world_centers = []

    for marker_id, marker_data in scene_markers.items():
        mid = int(marker_id)
        if mid not in robot_markers:
            continue

        for obs in marker_data.get("observations", []):
            if obs.get("camera_id") != camera_id:
                continue

            corners_px = np.array(obs["corners_px"], dtype=np.float32)
            rvec, tvec = solve_marker_pose(corners_px, K, dist, marker_size_m)
            if rvec is None:
                continue

            cam_centers.append(tvec.flatten())
            world_centers.append(robot_markers[mid])
            break

    if len(cam_centers) < 3:
        return None, None

    A = np.vstack(cam_centers)
    B = np.vstack(world_centers)
    R, t = rigid_transform_no_scale(A, B)
    return R, t


# ------------------------------------------------------------
# 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)
    print(f"[INFO] Loaded {len(robot_markers)} board markers from robot.json")

    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)

        R, t = build_camera_pose_from_board_markers(
            cam_id,
            scene["markers"],
            robot_markers,
            K,
            dist,
            args.marker_size
        )

        if R is None:
            print(f"[WARN] Camera {cam_id}: not enough board markers for pose estimation")
            continue

        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()