import os
import time
import h5py
import hydra
import numpy as np
from .recorder import Recorder
from openteach.utils.timer import FrequencyTimer
from openteach.constants import *
from openteach.utils.network import ZMQKeypointSubscriber, ZMQKeypointPublisher 


# To record robot information
class SimInformationRecord(Recorder):
    def __init__(
        self,
        port_configs,
        recorder_function_key,
        storage_path,
    ):

        # Data function and attributes
        #self.robot = hydra.utils.instantiate(robot_configs)
        #self.keypoint_function = self.robot.recorder_functions[recorder_function_key]
        self.robot = port_configs['robot']
        host = port_configs['host']
        jointanglesubscribeport = port_configs['jointanglesubscribeport']
        actualjointanglesubscribeport = port_configs['actualjointanglesubscribeport']
        timestampsubscribeport = port_configs['timestampssubscribeport']
        endeff_publish_port= port_configs['endeff_publish_port']
        endeffpossubscribeport= port_configs['endeffpossubscribeport']
        
        self.subscriber =  ZMQKeypointSubscriber(
            host = host,
            port = jointanglesubscribeport,
            topic= 'desired_angles'
        )

        self.jointanglesubscriber =  ZMQKeypointSubscriber(
            host = host,
            port = actualjointanglesubscribeport,
            topic= 'joint_states'
        )

        self.timestampsubscriber =  ZMQKeypointSubscriber(
            host = host,
            port = timestampsubscribeport,
            topic= 'timestamps'
        )

        self.endeffector_pos_subscriber =  ZMQKeypointSubscriber(
            host = host,
            port = jointanglesubscribeport,
            topic= 'endeffector_pos'
        )

        self.end_eff_coords_actual =  ZMQKeypointSubscriber(
            host = host,
            port = endeff_publish_port,
            topic= 'endeff_coords'
        ) 

        self.endeffector_pos_subscriber =  ZMQKeypointSubscriber(
                host = host,
                port = endeffpossubscribeport,
                topic= 'endeff_coords'
            )
        
    

        #For End effector 
        # Timer
        self.timer = FrequencyTimer(VR_FREQ)
        self.timestampsubscribeport= timestampsubscribeport
        self.recorder_function_key = recorder_function_key
        # Storage path for file
        self._filename = '{}_{}'.format(self.robot, recorder_function_key)
        self.notify_component_start('{}'.format(self._filename))
        self._recorder_file_name = os.path.join(storage_path, self._filename + '.h5')

        # Initializing data containers
        self.robot_information = dict()

    def stream(self):
        print('Checking if the keypoint port is active...')
        #while self.keypoint_function() is None:
            #continue

        print('Starting to record keypoints to store in {}.'.format(self._recorder_file_name))

        self.num_datapoints = 0
        self.record_start_time = time.time()

        while True:
            self.timer.start_loop()
            try:
                timestamps= self.timestampsubscriber.recv_keypoints()
                if self.robot=='robot_hand':
                    if self.recorder_function_key=='joint_states':
                        actual_joint_angles=self.jointanglesubscriber.recv_keypoints()
                    else:
                        actual_joint_angles=self.subscriber.recv_keypoints()
                elif self.robot=='moving_hand':
                    if self.recorder_function_key=='joint_states':
                        actual_joint_angles=self.jointanglesubscriber.recv_keypoints()
                    elif self.recorder_function_key=='cartesian_states':
                        actual_endeff_coords=self.end_eff_coords_actual.recv_keypoints()
                    elif self.recorder_function_key=='commanded_cartesian_states':
                        actual_endeff_coords=self.endeffector_pos_subscriber.recv_keypoints()
                    else:
                        actual_joint_angles=self.subscriber.recv_keypoints()
                else:
                    if self.recorder_function_key=='joint_states':
                        actual_endeff_coords=self.end_eff_coords_actual.recv_keypoints()
                    else:
                        actual_endeff_coords=self.endeffector_pos_subscriber.recv_keypoints()
               
                timestamps= self.timestampsubscriber.recv_keypoints()
                # proprio
                for key in ['position', 'timestamps']:  
                    if self.robot=='robot_hand':
                        if key not in self.robot_information.keys():
                            if key =='timestamps':
                                self.robot_information[key]=[timestamps]
                            elif key =='position':
                                self.robot_information[key]=[actual_joint_angles]   
                        else:
                            if key =='timestamps':
                                self.robot_information[key].append(timestamps)
                            elif key =='position':
                                self.robot_information[key].append(actual_joint_angles)
                    elif self.robot=='moving_hand':
                        if self.recorder_function_key == 'joint_states':
                            if key not in self.robot_information.keys():
                                if key =='timestamps':
                                    self.robot_information[key]=[timestamps]
                                elif key =='position':
                                    self.robot_information[key]=[actual_joint_angles]   
                            else:
                                if key =='timestamps':
                                    self.robot_information[key].append(timestamps)
                                elif key =='position':
                                    self.robot_information[key].append(actual_joint_angles)
                        else:
                            if key not in self.robot_information.keys():
                                if key =='timestamps':
                                    self.robot_information[key]=[timestamps]
                                elif key =='position':
                                    self.robot_information[key]=[actual_endeff_coords]   

                            else:
                                if key =='timestamps':
                                    self.robot_information[key].append(timestamps)
                                elif key =='position':
                                    self.robot_information[key].append(actual_endeff_coords)
                    else:
                        if key not in self.robot_information.keys():
                            if key =='timestamps':
                                self.robot_information[key]=[timestamps]
                            elif key =='position':
                                self.robot_information[key]=[actual_endeff_coords]   
                        else:
                            if key =='timestamps':
                                self.robot_information[key].append(timestamps)
                            elif key =='position':
                                self.robot_information[key].append(actual_endeff_coords)
                            


                self.num_datapoints += 1
                self.timer.end_loop()
            except KeyboardInterrupt:
                self.record_end_time = time.time()
                break

        # Displaying statistics
        self._display_statistics(self.num_datapoints)
        
        # Saving the metadata
        self._add_metadata(self.num_datapoints)

        # Writing to dataset
        print('Compressing keypoint data...')
        with h5py.File(self._recorder_file_name, "w") as file:
            # Main data
            #print(self.robot_information)
            print("Data Saved")
            for key in self.robot_information.keys():
                if key != 'timestamps':
                    self.robot_information[key] = np.array(self.robot_information[key], dtype = np.float32)
                    file.create_dataset(key +'s', data = self.robot_information[key], compression="gzip", compression_opts = 6)
                    print("file created")
                else:
                    self.robot_information['timestamps'] = np.array(self.robot_information['timestamps'], dtype = np.float64)
                    file.create_dataset('timestamps', data = self.robot_information['timestamps'], compression="gzip", compression_opts = 6)

            # Other metadata
            file.update(self.metadata)
        print('Saved keypoint data in {}.'.format(self._recorder_file_name))