#!/usr/bin/env python3
"""Recompute hand-eye calibration from an observations .npz file."""

import argparse
from typing import List

import cv2
import numpy as np


def _method_from_name(name: str) -> int:
    name = name.upper()
    methods = {
        "TSAI": cv2.CALIB_HAND_EYE_TSAI,
        "PARK": cv2.CALIB_HAND_EYE_PARK,
        "HORAUD": cv2.CALIB_HAND_EYE_HORAUD,
        "ANDREFF": cv2.CALIB_HAND_EYE_ANDREFF,
        "DANIILIDIS": cv2.CALIB_HAND_EYE_DANIILIDIS,
    }
    return methods.get(name, cv2.CALIB_HAND_EYE_TSAI)


def main() -> None:
    parser = argparse.ArgumentParser(description="Recompute hand-eye from observations npz.")
    parser.add_argument("observations", help="*_observations.npz path")
    parser.add_argument("--method", default="TSAI")
    args = parser.parse_args()

    data = np.load(args.observations, allow_pickle=True)
    R_gripper2base: List[np.ndarray] = list(data["R_gripper2base"])
    t_gripper2base: List[np.ndarray] = list(data["t_gripper2base"])
    R_target2cam: List[np.ndarray] = list(data["R_target2cam"])
    t_target2cam: List[np.ndarray] = list(data["t_target2cam"])

    method = _method_from_name(args.method)
    R_cam2gripper, t_cam2gripper = cv2.calibrateHandEye(
        R_gripper2base, t_gripper2base, R_target2cam, t_target2cam, method=method
    )

    T_cam2gripper = np.eye(4)
    T_cam2gripper[:3, :3] = R_cam2gripper
    T_cam2gripper[:3, 3] = t_cam2gripper.reshape(3)

    base_targets = []
    for i in range(len(R_gripper2base)):
        base_T_gripper = np.eye(4)
        base_T_gripper[:3, :3] = R_gripper2base[i]
        base_T_gripper[:3, 3] = np.array(t_gripper2base[i]).reshape(3)
        cam_T_target = np.eye(4)
        cam_T_target[:3, :3] = R_target2cam[i]
        cam_T_target[:3, 3] = np.array(t_target2cam[i]).reshape(3)
        base_T_target = base_T_gripper @ T_cam2gripper @ cam_T_target
        base_targets.append(base_T_target[:3, 3])
    base_targets = np.array(base_targets)
    mean_target = base_targets.mean(axis=0)
    errors = np.linalg.norm(base_targets - mean_target[None, :], axis=1)

    print("T_cam2gripper:")
    print(T_cam2gripper)
    print(f"Mean translation error: {errors.mean()*1000:.2f} mm")
    print(f"Max translation error: {errors.max()*1000:.2f} mm")


if __name__ == "__main__":
    main()