import os
os.environ["GLOG_minloglevel"] = "1"

import sys
from utils import root_path
sys.path.append(os.path.join(root_path, "cmake-build-debug/example"))
from robot_wrapper import Robot

import time
import signal
import numpy as np
from scipy.spatial.transform import Rotation as R
from utils import solve_ik
from utils import get_pose
from target_position import init_realsense, get_closest_red_point


def signal_handler(sig, frame):
    print("\n收到 Ctrl+C，准备退出...")
    raise SystemExit

signal.signal(signal.SIGINT, signal_handler)

def rt_to_transform(R, t):
    """拼接旋转矩阵和平移向量为齐次矩阵"""
    T = np.eye(4)
    T[:3, :3] = R
    T[:3, 3] = np.squeeze(t)
    return T

def transform_point(T, P):
    """用齐次变换矩阵 T (4x4) 把点 P(3,) 转换到新坐标系"""
    P_h = np.append(P, 1)   # [x, y, z, 1]
    P_new = T @ P_h
    return P_new[:3]

def compute_point_in_base(T_base_ee, R_ee_cam, t_ee_cam, P_cam_target):
    T_ee_cam = rt_to_transform(R_ee_cam, t_ee_cam)
    T_base_cam = T_base_ee @ T_ee_cam
    P_base_target = transform_point(T_base_cam, P_cam_target)
    return P_base_target

IDEL_Q = [-0.05804541534555635, 1.460164607187404, 1.458934384971377, 0.29042831678163805,
          -1.498103103566999, 0.039690803641003906, -0.08653132466712826]
INIT_POSE = [0.45, -0.25, -0.15, 0, 0, 0]

# FOLLOW_OFFSET = [-0.3, 0, 0]
# EE_LINK = "EE"
# FOLLOW_OFFSET = [-0.2, 0, 0]
# EE_LINK = "CAM"
FOLLOW_OFFSET = [-0.01, 0, 0]
EE_LINK = "FINGER"

def main():
    robot = Robot(os.path.join(root_path, "config/cabin_robot.xml"), "hc01")
    current_q = robot.getJointQ('right')
    diff = np.array(IDEL_Q) - np.array(current_q)
    if np.max(np.abs(diff)) >= 0.01:
        robot.moveJ("right", IDEL_Q)

    pipeline, align = init_realsense()
    ok, _, init_rq = solve_ik(INIT_POSE, base_link="PELVIS_S", target_link="R_FINGER_TIP")
    if not ok:
        print("solve ik init_pose failed, init_pose:", INIT_POSE)
        return
    print("init q:", init_rq)
    print("Robot movej: ", init_rq)
    robot.moveJ("right", init_rq)
    try:
        while True:
            current_q = robot.getJointQ('right')
            T_base_ee = get_pose(current_q, base_link="PELVIS_S", target_link="R_WRIST_R_S")
            T_base_cam = get_pose(current_q, base_link="PELVIS_S", target_link="R_CAM")
            T_base_tool = get_pose(current_q, base_link="PELVIS_S", target_link="R_FINGER_TIP")
            print(T_base_cam)
            P_cam_target = get_closest_red_point(pipeline, align, vis_path="red_point.png")
            if P_cam_target is None:
                print("can not find red point")
                time.sleep(2)
                continue

            T_cam_target = np.eye(4)
            T_cam_target[:3, 3] = P_cam_target
            T_base_target = T_base_cam @ T_cam_target
            P_base_target = T_base_target[:3, 3]
            desire_P_base_tool = P_base_target + FOLLOW_OFFSET

            if EE_LINK == "EE":
                desire_rpy_base_tool = R.from_matrix(T_base_ee[:3, :3]).as_euler('xyz', degrees=True)
                desire_pose = list(desire_P_base_tool) + list(desire_rpy_base_tool)
                ok, _, rq = solve_ik(desire_pose, base_link="PELVIS_S", target_link="R_WRIST_R_S")
            elif EE_LINK == "CAM":
                desire_rpy_base_tool = R.from_matrix(T_base_cam[:3, :3]).as_euler('xyz', degrees=True)
                desire_pose = list(desire_P_base_tool) + list(desire_rpy_base_tool)
                ok, _, rq = solve_ik(desire_pose, base_link="PELVIS_S", target_link="R_CAM")
            elif EE_LINK == "FINGER":
                desire_rpy_base_tool = R.from_matrix(T_base_tool[:3, :3]).as_euler('xyz', degrees=True)
                desire_pose = list(desire_P_base_tool) + list(desire_rpy_base_tool)
                ok, _, rq = solve_ik(desire_pose, base_link="PELVIS_S", target_link="R_FINGER_TIP")
            else:
                raise RuntimeError("undefined EE_LINK: ", EE_LINK)

            if not ok:
                print("solve ik failed, pose=", desire_pose)
                time.sleep(1)
            else:
                current_q = robot.getJointQ('right')
                diff = np.array(rq) - np.array(current_q)
                if np.max(diff) < 0.01:
                    continue
                print("move position:", desire_P_base_tool)
                print("Robot movej: ", rq)
                input_data = input("move robot? type \"yes\" to move, or type \"no\" to skip \n")
                if input_data == "yes":
                    robot.moveJ("right", rq)
                    time.sleep(5)
                    robot.moveJ("right", init_rq)

            print("Return to init pose, Robot movej: ", init_rq)

    except SystemExit:
        print("安全退出程序...")

    finally:
        robot.moveJ("right", IDEL_Q)
        pipeline.stop()


if __name__ == '__main__':
    main()