from hardware.dobot_nova5 import dobot_nova5
import cv2
from hardware.remote_cam import ToolCamera
import threading

import numpy as np
import os
class Getpose:
    def __init__(self):
        self.arm = dobot_nova5("192.168.5.1")
        self.arm.start()
        self.arm.init()
        self.cam = ToolCamera(host='127.0.0.1')
        self.cam.start()
        
        self.thread1=threading.Thread(target=self.show)
        self.thread1.start()
        self.pose_data=[]
        self.image_data=[]
        self.arm.start_drag()
    
    def show(self):
        while True:
            self.rgb, self.depth, intrinsics = self.cam.get_latest_frame()
            cv2.imshow(self.rgb)

    def add_data(self):
        angle=self.arm.getAngel()
        self.pose_data=self.pose_data.append(angle)
        print(angle)

    def save_data(self):
        data=np.array(self.pose_data)
        np.savetxt('pose.txt', data ,fmt='%.5f')

if __name__ == "__main__":
    save_directory = '/home/jsfb/jsfb_ws/global_config/captured_images'
    if not os.path.exists(save_directory):
        os.makedirs(save_directory)
    sele=Getpose()
    
    key = cv2.waitKey(1) & 0xFF
    while True:
        i=1
        if key == ord('s'):
            img_path = f'photo{i}.png'
            filename = os.path.join(save_directory, img_path)
            cv2.imwrite(filename, sele.rgb) # 保存照片
            sele.add_data()
            
            filename2 = os.path.join(save_directory, f'pose.txt')
            np.savetxt(filename2, sele.pose_data)
            # pose = increase_dof(poselist[index])
            
            i += 1
            # 按下'q'键退出循环
        elif key == ord('q'):
            sele.arm.disableRobot()
            break