from openteach.constants import VR_FREQ,  ARM_LOW_RESOLUTION, ARM_HIGH_RESOLUTION, ARM_TELEOP_CONT, ARM_TELEOP_STOP , GRIPPER_CLOSE, GRIPPER_OPEN
from openteach.components import Component
from openteach.utils.timer import FrequencyTimer
from openteach.utils.network import create_pull_socket, ZMQKeypointPublisher
            
class OculusVRTwoHandDetector(Component):
    def __init__(self, 
                 host, 
                 oculus_right_port,
                 oculus_left_port,
                 keypoint_pub_port, 
                 button_port, 
                 button_publish_port
    ):
        self.notify_component_start('vr detector')
        # Initializing the network socket
        self.raw_keypoint_right_socket = create_pull_socket(host, oculus_right_port)
        self.button_keypoint_socket = create_pull_socket(host, button_port)
        #self.teleop_reset_socket = create_pull_socket(host, teleop_reset_port)
        self.raw_keypoint_left_socket= create_pull_socket(host, oculus_left_port)
        # ZMQ Keypoint publisher
        self.hand_keypoint_publisher = ZMQKeypointPublisher(
            host = host,
            port = keypoint_pub_port
        )
        # button_publish_port = 8093
        self.button_socket_publisher = ZMQKeypointPublisher(
            host = host,
            port = button_publish_port
        ) 

        self.timer = FrequencyTimer(VR_FREQ)

    def _process_data_token(self, data_token):
        return data_token.decode().strip()

    def _extract_data_from_token(self, token):        
        data = self._process_data_token(token)
        # information = dict(hand = 'right' if data.startswith('right') else 'left')
        information = dict()
        keypoint_vals = [0] if data.startswith('absolute') else [1]

        # Data is in the format <hand>:x,y,z|x,y,z|x,y,z
        vector_strings = data.split(':')[1].strip().split('|')
        for vector_str in vector_strings:
            vector_vals = vector_str.split(',')
            for float_str in vector_vals[:3]:
                keypoint_vals.append(float(float_str))
            
        information['keypoints'] = keypoint_vals
        return information

    def _publish_right_data(self, keypoint_dict):
        self.hand_keypoint_publisher.pub_keypoints(
            keypoint_array = keypoint_dict['keypoints'], 
            topic_name = 'right'
        )

    def _publish_left_data(self, keypoint_dict):
        self.hand_keypoint_publisher.pub_keypoints(
            keypoint_array= keypoint_dict['keypoints'],
            topic_name= 'left'
        )

    def _publish_button_data(self,button_feedback):
        self.button_socket_publisher.pub_keypoints(
            keypoint_array = button_feedback, 
            topic_name = 'button'
        )

    # def _publish_pause_data(self,pause_status):
    #     self.teleop_reset_publisher.pub_keypoints(
    #         keypoint_array = pause_status, 
    #         topic_name = 'pause'
    #     )

    

    def _publish_gripper_port(self, gripper_left_status, gripper_right_status):
        self.gripper_port_publisher.pub_keypoints(
            keypoint_array= gripper_left_status,
            topic_name = 'gripper_left' 
        )
        self.gripper_port_publisher.pub_keypoints(
            keypoint_array= gripper_right_status,
            topic_name = 'gripper_right'
        )

    def _publish_pause_left_right_data(self,pause_left_status, pause_right_status):
        self.pause_port_publisher.pub_keypoints(
            keypoint_array = pause_left_status, 
            topic_name = 'pause_left'
        )
        self.pause_port_publisher.pub_keypoints(
            keypoint_array= pause_right_status,
            topic_name= 'pause_right'
        )

    

    def stream(self):

        while True:
            # try:
            self.timer.start_loop()

            # Getting the raw keypoints
            raw_right_keypoints = self.raw_keypoint_right_socket.recv()
            print("Right raw Keypoints", raw_right_keypoints)
            raw_left_keypoints = self.raw_keypoint_left_socket.recv()
            print("Left raw Keypoints", raw_left_keypoints)
            button_feedback = self.button_keypoint_socket.recv()
            print("Button Feedback", button_feedback)

            if button_feedback==b'Low':
                button_feedback_num = ARM_LOW_RESOLUTION
            else:
                button_feedback_num = ARM_HIGH_RESOLUTION

            # Processing the keypoints and publishing them
            keypoint_right_dict = self._extract_data_from_token(raw_right_keypoints)
            keypoint_left_dict = self._extract_data_from_token(raw_left_keypoints)
            self._publish_right_data(keypoint_right_dict)
            print("Published Right Keypoints")
            self._publish_left_data(keypoint_left_dict)
            print("Published Left Keypoints")
            self._publish_button_data(button_feedback_num)
            print("Published Button Feedback")
            #self._publish_pause_data(pause_status)
            #self._publish_gripper_port(gripper_left_status,gripper_right_status)
            #self._publish_pause_left_right_data(pause_left_status,pause_right_status)
            # self._publish_gripper_rotate_data(gripper_left_status,gripper_right_status)

            #print("Oculus.py working fine")
            
            self.timer.end_loop()

            # except KeyboardInterrupt:
            #     break

        self.raw_keypoint_socket.close()
        self.hand_keypoint_publisher.stop()
        print('Stopping the oculus keypoint extraction process.')