import os
import sys
import cv2
import threading
import numpy as np
import time

current_file_path = os.path.abspath(__file__)
dobot_nova5_path = os.path.dirname(os.path.dirname(current_file_path))
sys.path.append(dobot_nova5_path)

from hardware.dobot_nova5 import dobot_nova5
from hardware.remote_cam import ToolCamera

class Getpose:
    def __init__(self):
        self.running = True
        self.arm = dobot_nova5()
        self.arm.start()
        self.arm.init()
        self.cam = ToolCamera(host='127.0.0.1')
        self.cam.start()
        time.sleep(2)
        
        self.pose_data = []
        self.image_data = []
        # self.arm.start_drag()
        self.rgb = None
            
    def test_show(self):
        self.rgb, self.depth, intrinsics = self.cam.get_latest_frame()
        if self.rgb is not None:
            cv2.imshow("Video Feed", self.rgb)
            cv2.waitKey(1000)
    
    def add_data(self):
        angle = self.arm.getAngel()
        self.pose_data.append(angle)
        print(f"Current angle: {angle}")
    
    def save_data(self):
        if self.pose_data:
            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)
    
    myTest = Getpose()
    
    myTest.arm.start_drag()
    
    time.sleep(0.5)
    try:
        count = 1
        i = 1
        while True:
            # 读取摄像头的帧
            r, d, intrinsics = myTest.cam.get_latest_frame()
            # ret, frame = cam.read()
            frame = r

            # if not ret:
            #     print("Error: Could not read frame from video device.")
            #     break

                        # 显示摄像头的帧
            cv2.imshow('found', frame)

            # 等待按键事件
            key = cv2.waitKey(1) & 0xFF


            # if key == ord('r'):
            #     pose = increase_dof(poselist[index])
            #     index+=1
            #     code = cp.arm.move_joint(pose, 4,wait=True)
            #     if code == -1:
            #         print("运动到拍照位置失败")
            #     else:

            # 按下's'键保存照片
            if key == ord('s'):
                img_path = f'photo{i}'
                if not os.path.exists(os.path.join(save_directory, img_path)):
                    os.makedirs(os.path.join(save_directory, img_path))
                filename = os.path.join(save_directory, f'{img_path}/rgb_image.png')
                cv2.imwrite(filename, frame) # 保存照片
                filename1 = os.path.join(save_directory, f'{img_path}/depth_image.png')
                cv2.imwrite(filename1, d) # 保存照片
                angle = myTest.arm.Get_feedInfo_angle()
                print("angle = ",angle)
                filename2 = os.path.join(save_directory, f'{img_path}/angle.txt')
                np.savetxt(filename2, angle.shape(-1,6),delimiter=',')
                print(f'Saved {filename}')
                
                # pose = increase_dof(poselist[index])
                # index+=1
                # code = cp.arm.move_joint(poselist[index], 4,wait=True)
                # if code == -1:
                #     print("运动到拍照位置失败")
                count += 1 # 增加照片计数
                i += 1
            # 按下'q'键退出循环
            elif key == ord('q'):
                break
            

    except KeyboardInterrupt:
        # 处理Ctrl+C中断
        print("\nCamera session ended by user.")
    finally:
        # 释放资源
        myTest.arm.stop_drag()
        time.sleep(0.5)
        myTest.arm.disableRobot()
        myTest.cam.stop()
        cv2.destroyAllWindows()

        # 总结拍摄的照片
        print(f'Total number of images captured: {count}')