# AUTOGENERATED! DO NOT EDIT! File to edit: 13_PokingBot_training_multi_machine.ipynb (unless otherwise specified).

__all__ = ['CURL_NAVI_GYM_V1']

# Cell
import socket
import _thread
import pickle
import time
import os
import sys
import signal
import subprocess
import numpy as np
import yaml
import random
import gym
from gym import error, spaces, utils
from gym.utils import seeding
import math
from math import e

# Cell

from .CurriculumSchedulerMultiMachine import CurriculumSchedulerMultiMachine
class CURL_NAVI_GYM_V1(gym.Env):
    def __init__(self):
        with open('Data/IP_ENV.yaml', 'r') as f: self.IP_ENV = yaml.load(f, Loader=yaml.CLoader)
        # Spaces
        self.observation_space = spaces.Box(high=100, low=0, shape=( (4+4)*241+10*3, ) ) # laser max range=100
        self.action_space = spaces.Box(low=np.array([0, -1, 0]), high=np.array([1, 1, 1]), dtype=np.float32)
        # state info
        self.info = {'laser_shape': (4+4, 241),'goal_shape': (10, 3)} # 4 frames laser range, 4 frames laser label, 10 frames pose

        self.observation_state = np.zeros(((4+4)*241+10*3,), dtype=np.float32)
        self.action_state = np.array([0,0,0], dtype=np.float32)
        self.reward = float(0)
        self.done_state = (False,False) # done & collision -> (True,False),
                                        # done & goal reached -> (True,True)
        self.reset_gazebo_model = False
        self.start_goal_pair = {'goal': [0,0,1], 'start': [0,0,1]}
        self.action_scale = {'angular': 0.8, 'linear': 1.5, 'ur5': 0.5}
        self.have_received_data = False
        self.step_count = 0

        self.scheduler = CurriculumSchedulerMultiMachine('Data/curriculum.yaml')
        self.socket = socket.socket()         # Create a socket object
        host = '0.0.0.0' # Get local machine name
        port = 9453      # Reserve a port for your service.
        print('Server started!')
        print('Waiting for clients...')

        self.socket.bind((host, port))        # Bind to the port
        self.socket.listen(15)                 # Now wait for client connection.
        _thread.start_new_thread(self.socket_server_thread,(1,))


    def socket_server_thread(self,_):
        while True:
            c, addr = self.socket.accept()     # Establish connection with client
            _thread.start_new_thread(self.on_new_client,(c,addr))

    def on_new_client(self, clientsocket, addr):
        while True:
            env = self.IP_ENV[addr[0]]

            if not (env == self.scheduler.current_task_env):
                msg = clientsocket.recv(256000)
                pkg = {}
                clientsocket.send(pickle.dumps(pkg))
                continue

            # receive
            time.sleep(0.2) # sleeping is very important !!!
            msg = clientsocket.recv(256000)
            try:
                pack_from_client = pickle.loads(msg)
                self.observation_state = pack_from_client["observation_state"]
                self.reward = pack_from_client["reward"]
                self.done_state = pack_from_client["done_state"]
                self.have_received_data = True
            except Exception as e:
                pass

            #send
            pack_to_client = { "action_state": self.action_state, \
                               "reset_gazebo_model": self.reset_gazebo_model, \
                               "start_goal_pair": self.start_goal_pair , \
                               "action_scale": self.action_scale }
            clientsocket.send(pickle.dumps(pack_to_client))
            if self.reset_gazebo_model : self.reset_gazebo_model = False

        clientsocket.close()


    def reset(self):
        self.start_goal_pair, self.action_scale = self.scheduler.update(self.done_state[1])
        self.reset_gazebo_model = True
        self.step_count = 0
        return self.observation_state

    def step(self, action):
        while not self.have_received_data:
            pass
        self.action_state = action
        self.step_count += 1
        self.have_received_data = False
        return self.observation_state, self.reward, self.done_state[0], self.info

