# utils_robot.py
# 同济子豪兄 2024-5-22
# 启动并连接机械臂，导入各种工具包

print('导入机械臂连接模块')

from lebai import zeroconf, l_master
import cv2
import numpy as np
import time

# 连接机械臂
# 实例化一个Discovery类用来存储设备发现的信息
d = zeroconf.Discovery()
# 搜索局域网内设备，返回一个最多16个元素的列表
controllers = d.resolve()
robot = l_master.Robot(controllers[0].ip_address)

def back_zero():
    '''
    机械臂归零
    '''
    print('机械臂归零')
    robot.movej(
        {"j1": 0.0,
         "j2": 0.0,
         "j3": 0.0,
         "j4": 0.0,
         "j5": 0.0,
         "j6": 0.0},
        1.0, 1.0, 0.0, 0.0)
    time.sleep(3)

def relax_arms():
    print('放松机械臂关节')
    robot.teach_mode()

def move_to_coords(X=150, Y=-130, Z=0.35):
    print('移动至指定坐标：X {} Y {}'.format(X, Y))
    X, Y = eye2hand(X, Y)
    cartesian_pose = {'x' : X, 'y' : Y, 'z' : Z, 'rz' : 0, 'ry' : 0, 'rx' : np.pi} # 目标位姿笛卡尔数据
    a = 0.5 # 关节加速度 (rad/s2)
    v = 0.2 # 关节速度 (rad/s)
    t = 0   # 运动时间 (s)。 当 t > 0 时，参数速度 v 和加速度 a 无效
    r = 0   # 交融半径 (m)。用于指定路径的平滑效果
    robot.movel(cartesian_pose, a, v, t, r) # 直线运动 https://help.lebai.ltd/sdk/motion.html#%E7%9B%B4%E7%BA%BF%E8%BF%90%E5%8A%A8
    robot.wait_move() # 等待运动完成
    time.sleep(4)

def single_joint_move(joint_index, angle):
    print('关节 {} 旋转至 {} 度'.format(joint_index, angle))
    robot.movej(
        {
            f"j{joint_index}" : angle / 180 * np.pi
        },
        1.0, 1.0, 0.0, 0.0
    )
    time.sleep(2)

def move_to_top_view():
    print('移动至俯视姿态')
    joint_pose = [np.pi / 2, -np.pi / 2, -np.pi / 2, -np.pi / 2, np.pi / 2, 0] # 目标位姿关节数据
    a = 0.5 # 关节加速度 (rad/s2)
    v = 0.2 # 关节速度 (rad/s)
    t = 0   # 运动时间 (s)。 当 t > 0 时，参数速度 v 和加速度 a 无效
    r = 0   # 交融半径 (m)。用于指定路径的平滑效果
    robot.movej(joint_pose, a, v, t, r) # 关节运动 https://help.lebai.ltd/sdk/motion.html#%E5%85%B3%E8%8A%82%E8%BF%90%E5%8A%A8

def top_view_shot(check=False):
    '''
    拍摄一张图片并保存
    check：是否需要人工看屏幕确认拍照成功，再在键盘上按q键确认继续
    '''
    print('    移动至俯视姿态')
    move_to_top_view()
    
    # 获取摄像头，传入0表示获取系统默认摄像头
    cap = cv2.VideoCapture(0, cv2.CAP_DSHOW)
    # 打开cap
    cap.open(0, cv2.CAP_DSHOW)
    time.sleep(0.3)
    success, img_bgr = cap.read()
    
    # 保存图像
    print('    保存至temp/vl_now.jpg')
    cv2.imwrite('temp/vl_now.jpg', img_bgr)

    # 屏幕上展示图像
    cv2.destroyAllWindows()   # 关闭所有opencv窗口
    cv2.imshow('vlm_agent', img_bgr) 
    
    if check:
        print('请确认拍照成功，按c键继续，按q键退出')
        while(True):
            key = cv2.waitKey(10) & 0xFF
            if key == ord('c'): # 按c键继续
                break
            if key == ord('q'): # 按q键退出
                # exit()
                cv2.destroyAllWindows()   # 关闭所有opencv窗口
                raise NameError('按q退出')
    else:
        if cv2.waitKey(10) & 0xFF == None:
            pass
        
    # 关闭摄像头
    cap.release()
    # 关闭图像窗口
    # cv2.destroyAllWindows()

def eye2hand(X_im=160, Y_im=120):
    '''
    输入目标点在图像中的像素坐标，转换为机械臂坐标
    '''
    cali_ims = [
        (270, 223),
        (592, 279),
        (418, 333),
    ]
    
    cali_mcs = [
        (0.190, 0.229),
        (-0.095, 0.270),
        (0.065, 0.300),
    ]
    
    
    X_cali_im = [cali_im[0] for cali_im in cali_ims]  # 像素坐标
    X_cali_mc = [cali_mc[0] for cali_mc in cali_mcs]  # 机械臂坐标
    Y_cali_im = [cali_im[1] for cali_im in cali_ims]  # 像素坐标
    Y_cali_mc = [cali_mc[1] for cali_mc in cali_mcs]  # 机械臂坐标

   # 拟合三次多项式
    X_poly = np.polyfit(X_cali_im, X_cali_mc, deg=2)
    Y_poly = np.polyfit(Y_cali_im, Y_cali_mc, deg=2)

    # 使用多项式模型进行插值
    X_mc = np.polyval(X_poly, X_im)
    Y_mc = np.polyval(Y_poly, Y_im)

    return X_mc, Y_mc

def gripper_open():
    '''
    打开夹爪
    '''
    robot.set_claw(50, 100)
    
def gripper_close():
    '''
    关闭夹爪
    '''
    robot.set_claw(10, 0)

# 夹爪抓取并移动物体
def gripper_move(XY_START=[0.383, 0.121], HEIGHT_START=0.20, XY_END=[0.3, 0.12], HEIGHT_END=0.20, HEIGHT_SAFE=0.35):

    '''
    用夹爪，将物体从起点吸取移动至终点

    XY_START：起点机械臂坐标 (m)
    HEIGHT_START：起点高度 (m)，方块用 0.09，药盒子用 0.07
    XY_END：终点机械臂坐标 (m)
    HEIGHT_END：终点高度 (m)
    HEIGHT_SAFE：搬运途中安全高度 (m)
    '''
    
    a = 0.5 # 关节加速度 (rad/s2)
    v = 0.2 # 关节速度 (rad/s)
    t = 0   # 运动时间 (s)。 当 t > 0 时，参数速度 v 和加速度 a 无效
    r = 0   # 交融半径 (m)。用于指定路径的平滑效果
    
    # 夹爪移动至物体上方
    print('    夹爪移动至物体上方')
    cartesian_pose = {'x' : XY_START[0], 'y' : XY_START[1], 'z' : HEIGHT_SAFE, 'rz' : 0, 'ry' : 0, 'rx' : np.pi} # 目标位姿笛卡尔数据
    robot.movel(cartesian_pose, a, v, t, r) # 直线运动 https://help.lebai.ltd/sdk/motion.html#%E7%9B%B4%E7%BA%BF%E8%BF%90%E5%8A%A8
    robot.wait_move() # 等待运动完成

    # 夹爪向下抓取物体
    print('    夹爪向下抓取物体')
    gripper_open()
    cartesian_pose = {'x' : XY_START[0], 'y' : XY_START[1], 'z' : HEIGHT_START, 'rz' : 0, 'ry' : 0, 'rx' : np.pi}
    robot.movel(cartesian_pose, a, v, t, r)
    robot.wait_move()
    gripper_close()

    # 升起物体
    print('    升起物体')
    cartesian_pose = {'x' : XY_START[0], 'y' : XY_START[1], 'z' : HEIGHT_SAFE, 'rz' : 0, 'ry' : 0, 'rx' : np.pi}
    robot.movel(cartesian_pose, a, v, t, r)
    robot.wait_move()

    # 搬运物体至目标上方
    print('    搬运物体至目标上方')
    cartesian_pose = {'x' : XY_END[0], 'y' : XY_END[1], 'z' : HEIGHT_SAFE, 'rz' : 0, 'ry' : 0, 'rx' : np.pi}
    robot.movel(cartesian_pose, a, v, t, r)
    robot.wait_move()

    # 向下放下物体
    print('    向下放下物体')
    cartesian_pose = {'x' : XY_END[0], 'y' : XY_END[1], 'z' : HEIGHT_END, 'rz' : 0, 'ry' : 0, 'rx' : np.pi}
    robot.movel(cartesian_pose, a, v, t, r)
    robot.wait_move()
    gripper_open()
    time.sleep(1)

    # 移动至俯视姿态
    print('    移动至俯视姿态')
    move_to_top_view()