import hydra
import pickle
from openteach.components import Component
from openteach.utils.vectorops import get_distance
from openteach.utils.network import create_response_socket
from openteach.utils.timer import FrequencyTimer
from openteach.utils.network import ZMQKeypointSubscriber, ZMQKeypointPublisher
from openteach.constants import DEPLOY_REACH_THRESHOLD ,DEPLOY_FREQ
from multiprocessing import Process

class DeployServerSim(Component):
    def __init__(self, configs):
        self.configs = configs
        
        # Initializing the camera subscribers
        self._init_robot_subscribers()
        
        self.deployment_socket = create_response_socket(
            host = self.configs.host_address,
            port = self.configs.deployment_port
        )

        #if self.configs.manual_sleep:
            #print('Using manual sleep.')
        self.timer = FrequencyTimer(DEPLOY_FREQ)
        self.endeff_velocity_publisher = ZMQKeypointPublisher(
            host = self.configs.host_address,
            port = self.configs.endeff_publish_port
        )

        self.joint_angle_subscriber = ZMQKeypointSubscriber(
            host = self.configs.host_address,
            port = self.jointanglesubscribeport,
            topic= 'current_angles'
        )
        # self.joint_angle_publisher = ZMQKeypointPublisher(
        #     host = self.configs.host_address,
        #     port = self.configs.jointangle_publish_port
        # )
        #else:
            #print('Using movement based sleep.')

    def _wait_for_movement(self, robot_action):
        print('Waiting to reach target state!')
        while True:
            achieved_target = 0
            #for robot in robot_action.keys():
                #current_state = self._robots[robot].get_joint_state()
            current_state = self.joint_angle_subscriber.recv_keypoints()
            #endeff_vel = self.
            target_action = robot_action[robot]

            distance = get_distance(current_state['position'], target_action)
            if distance < DEPLOY_REACH_THRESHOLD:
                achieved_target += 1

            if achieved_target == len(robot_action.keys()):
                break

        print('Achieved target state!')

    # def _init_robot_subscribers(self):
    #     robot_controllers = hydra.utils.instantiate(self.configs.robot.controllers)
    #     self._robots = dict()
    #     for robot in robot_controllers:
    #         self._robots[robot.name] = robot

    #def _init_sim_subscribers(self):
    
    def _init_sim_environment(self):
         for env_config in self.configs.robot.environment:
            self.processes.append(Process(
                target = self._start_component,
                args = (env_config, )
            ))

        

    def _perform_robot_action(self, robot_action_dict):
        #try:
            for robot in robot_action_dict.keys(): 
                if robot not in self._robots.keys():
                    print('Robot: {} is an illegal argument.'.format(robot))
                    return False
                else:
                    if robot == 'kinova': 
                        print("Moving kinova according to cartesian coordinates")
                        #self._robots[robot].move_coords(robot_action_dict[robot])
                        

                    else:
                        print("Moving allego according to joint angles")
                        self._robots[robot].move(robot_action_dict[robot])
                print('Applying action {} on robot: {}'.format(robot_action_dict[robot], robot))
            return True
        #except:
            #return False

    def _start_sim_operator_component(self, configs):
        
        component = hydra.utils.instantiate(configs)
        component.stream()

    def _get_robot_states(self):
        data = dict()
        for robot_name in self._robots.keys():
            if robot_name == 'kinova':
                data[robot_name] = self._robots[robot_name].get_cartesian_position() # Get the position only
                print(f'data[{robot_name}].shape: {data[robot_name].shape}')
            else: # allegro
                data[robot_name] = self._robots[robot_name].get_joint_state()

        return data

    def _send_robot_state(self):
        robot_states = self._get_robot_states()
        self.deployment_socket.send(pickle.dumps(robot_states, protocol = -1))

    def stream(self):
        self.notify_component_start('robot deployer')
        while True:
            #try:
                print('\nListening')
                robot_action = pickle.loads(self.deployment_socket.recv())

                if robot_action == 'get_state':
                    print('Requested for robot state information.')
                    self._send_robot_state()
                    continue

                success = self._perform_robot_action(robot_action)

                # Sleep to achieve the action
                #if self.configs.manual_sleep:
                self.timer.start_loop()
                self.timer.end_loop()
                #else:
                    #self._wait_for_movement(robot_action)

                if success:
                    self._send_robot_state()
                    print('Applied robot action.')
                else:
                    self.deployment_socket.send("Command failed!")
            #except:
                #print('Illegal values passed. Terminating session.')
                #break

        print('Closing robot deployer component.')
        self.deployment_socket.close()