import os
import numpy as np
import cv2
import sounddevice as sd
import pandas as pd
import gym
from ..env_bases import BaseEnv
import pybullet as p
from cfg import main_config
from ..robot_locomotors import TurtleBot
from Envs.audioLoader import audioLoader
from ..scene_abstract import SingleRobotEmptyScene


class RLEnvRSI2(BaseEnv):
	def __init__(self):

		self.config=main_config() # load env config
		self.audio = audioLoader(config=self.config)
		self.robot = TurtleBot(config=self.config)

		# observation space
		d = {
			'image': gym.spaces.Box(low=0, high=255, shape=self.config.img_dim, dtype='uint8'),
			'goal_sound': gym.spaces.Box(low=-np.inf, high=np.inf, shape=self.config.sound_dim, dtype=np.float32),
			'current_sound': gym.spaces.Box(low=-np.inf, high=np.inf, shape=self.config.sound_dim, dtype=np.float32),
			'robot_pose': gym.spaces.Box(low=-np.inf, high=np.inf,shape=(2,), dtype=np.float32),
			'ground_truth': gym.spaces.Box(low=0, high=4, shape=(1,), dtype=np.int32), # 0: cube, 1: sphere, 2:cone, 3: cylinder
		}

		self.observation_space = gym.spaces.Dict(d)
		self.maxSteps = self.config.RLEnvMaxSteps

		# setup action space
		high = np.ones(self.config.RLActionDim)
		self.action_space = gym.spaces.Box(-high, high, dtype=np.float32)

		BaseEnv.__init__(self, self.robot,
						 config=self.config,
						 render=self.config.render,
						 action_space=self.action_space,
						 observation_space=self.observation_space)

		# 3D models
		self.arena_floor = None
		self.arenaWallInner = None
		self.arenaWallOuter = None
		self.wallTexPath = os.path.join(self.config.mediaPath, 'arena', 'wallAppearance', 'texture')
		self.wallTextureList = []

		self.objTexDict = {}
		self.objTexPath = os.path.join(self.config.mediaPath, 'objects', 'texture')
		self.objTexList = []

		# all objects that the robot will need to reach.  e.g. ['cube','sphere','cone','cylinder']
		# cube: choice 0, sphere: choice 1, cone: choice 2, cylinder: choice 3
		# The information related to them will be given by the same order in this list.
		self.objList = self.config.objList
		self.objUidList = []  # PyBullet Uid for each object

		# each spot associates with a list [x,y,yaw]. The order will be same with self.objList
		# [[x,y,yaw],[x,y,yaw],...]
		self.objPoseList = []  # objects position and orientation information
		# current obj in the scene. e.g. [2,3] means cone and cylinder in scene
		self.objInScene = np.array([])

		self.goalObjIdx = None
		self.goal_sound = None
		self.ground_truth = None
		self.goal_audio = None

		# ray test
		self.rayIds = []
		self.collision = False  # does collision happen during an episode

		self.size_per_class = np.zeros((len(self.config.objList),), dtype=np.int64)
		for key in self.config.soundSource['size']:
			self.size_per_class = self.size_per_class + self.config.soundSource['size'][key]
		self.size_per_class_cumsum = np.cumsum(self.size_per_class)

		self.goal_area_count = 0

		if self.config.record:
			self.episodeRecorder = Recorder()  # used for recording data
			self.episodeRecorder.saveTo=self.config.recordSaveDir
			if self.config.loadAction:
				self.episodeRecorder.loadActions()
				self.episodeRecorder.loadFrom = self.config.loadActionFile

	def create_single_player_scene(self, bullet_client):
		"""
		Setup physics engine and simulation
		:param bullet_client:
		:return: a scene
		"""
		# use large frame skip such that motor has time to follow the desired position
		return SingleRobotEmptyScene(bullet_client, gravity=(0, 0, -9.8),
									 timestep=self.timeStep, frame_skip=self.config.frameSkip,
									 render=self.config.render)

	def loadTex(self):
		wallTexList = os.listdir(self.wallTexPath)
		idx = np.arange(len(wallTexList))
		self.np_random.shuffle(idx)

		# load texture for walls
		for i in range(self.config.numTexture):
			# key=fileName, val=textureID. If we have already loaded the texture, no need to reload and drain the memory
			texID = self._p.loadTexture(os.path.join(self.wallTexPath, wallTexList[idx[i]]))
			self.wallTextureList.append(texID)

		# load texture for objects
		objTextureList = os.listdir(self.objTexPath)
		for i in range(len(objTextureList)):
			texID = self._p.loadTexture(os.path.join(self.objTexPath, objTextureList[i]))
			self.objTexDict[objTextureList[i]] = texID
		self.objTexList = list(self.objTexDict.keys())

	def loadArena(self):
		if self.config.domainRandomization:
			# load texture
			self.loadTex()
			# load arena: outer wall
			self.arenaWallOuter = self._p.loadURDF(
				os.path.join(self.config.mediaPath, 'arena', 'arena_wall', 'arena_wall2.urdf'),
				[0.25, 0, -0.02],
				[0.000000, 0.000000, 0.0, 1.0], flags=self._p.URDF_USE_MATERIAL_COLORS_FROM_MTL)

		# load arena: inner wall
		self.arenaWallInner = self._p.loadURDF(
			os.path.join(self.config.mediaPath, 'arena', 'arena_wall', 'arena_wall.urdf'),
			[0, 0, -0.02], [0., 0., 0.0, 1.0],
			flags=self._p.URDF_USE_MATERIAL_COLORS_FROM_MTL)

		# add collision shapes around the arena
		collisionID = self._p.createCollisionShape(shapeType=self._p.GEOM_BOX,
												   halfExtents=[1.0, 0.05, 0.05],
												   collisionFramePosition=[0.0, 0.0, 0.0],
												   collisionFrameOrientation=self._p.getQuaternionFromEuler([0, 0, 0]))

		visualID = self._p.createVisualShape(shapeType=self._p.GEOM_BOX,
											 halfExtents=[1.0, 0.05, 0.05],
											 visualFramePosition=[0, 0, 0],
											 visualFrameOrientation=self._p.getQuaternionFromEuler([0, 0, 0]),
											 rgbaColor=[1, 1, 1, 0]
											 )
		posList = [[0.0, 0.95, 0.05], [0.0, -0.95, 0.05], [0.95, 0.0, 0.05], [-0.95, 0.0, 0.05]]
		ornList = [[0, 0, 0], [0, 0, 0], [0, 0, np.pi / 2], [0, 0, np.pi / 2]]
		for i in range(4):
			objID = self._p.createMultiBody(baseMass=0,
											baseInertialFramePosition=[0, 0, 0],
											baseCollisionShapeIndex=collisionID,
											baseVisualShapeIndex=visualID,
											basePosition=posList[i],
											baseOrientation=self._p.getQuaternionFromEuler(ornList[i]))
		# load arena: floor
		self.arena_floor = self._p.loadURDF(
			os.path.join(self.config.mediaPath, 'arena', 'arena_floor', 'arena_floor.urdf'),
			[0, 0, 0.03], [0., 0., 0.0, 1.0],
			flags=self._p.URDF_USE_MATERIAL_COLORS_FROM_MTL)

	def loadObj(self, path, name):

		visualID = self._p.createVisualShape(shapeType=self._p.GEOM_MESH,
											 fileName=path,
											 visualFramePosition=[0, 0, 0],
											 visualFrameOrientation=self._p.getQuaternionFromEuler([0, 0, 0]),
											 rgbaColor=[1, 1, 1, 1]
											 )
		# create collision shape for each object so that the robot will not get closer than
		# objectsRadius+objectsExpandDistance to the object. It will be physically stopped by this collision shape
		radius_inner = self.config.objectsRadius[name]
		radius_outer = self.config.objectsRadius[name] + self.config.objectsExpandDistance[name]
		if name == 'cylinder':
			z_offset = 0.0
		else:
			z_offset = -0.1
		collisionID = p.createCollisionShapeArray(shapeTypes=[p.GEOM_CYLINDER, p.GEOM_CYLINDER],
												  radii=[radius_inner, radius_outer],
												  lengths=[0.5, 0.3],
												  collisionFramePositions=[[0., 0., 0.], [0., 0., z_offset]]
												  )

		objID = self._p.createMultiBody(baseMass=0,
										baseInertialFramePosition=[0, 0, 0],
										baseCollisionShapeIndex=collisionID,
										baseVisualShapeIndex=visualID,
										basePosition=[3, 3, 3],
										baseOrientation=self._p.getQuaternionFromEuler([0, 0, 0]))
		return objID

	def randomization(self):
		self.objPoseList = []
		# Clean the arena
		for i in range(len(self.objList)):  # for each object
			self._p.resetBasePositionAndOrientation(self.objUidList[i], [0, 0, -1],
													self._p.getQuaternionFromEuler([0, 0, 0]))
		# plan
		currentPlan = {}

		feasiblePlan = False
		while not feasiblePlan:
			# robot
			robotInitRegion_radius = self.config.robotInitRegion_radius
			robotx, roboty, robotYaw = self.randomXYYaw(rouLow=0, rouHigh=robotInitRegion_radius, thetaLow=0,
														thetaHigh=2 * np.pi, yawLow=-np.pi, yawHigh=np.pi)

			# robot is initialized within robotInitRegion_radius. It will not hit the wall
			currentPlan['robot'] = [robotx, roboty, robotYaw]

			# deal with the objects
			for index, obj in enumerate(self.objList):  # for each object
				if index in self.objInScene:  # put into the arena
					collisionFree = False
					trial = 0
					while not collisionFree:
						trial = trial + 1
						if trial > 200:
							break
						objx, objy, objYaw = self.randomXYYaw(rouLow=robotInitRegion_radius, rouHigh=1, thetaLow=0,
															  thetaHigh=2 * np.pi, yawLow=-np.pi, yawHigh=np.pi)
						new_obj = (obj, [objx, objy, objYaw],)
						collide = self.reject_fn(new_obj, currentPlan)
						if not collide:
							currentPlan[obj] = [objx, objy, objYaw]
							collisionFree = True

					if not collisionFree:
						break
				else:  # hide under the arena
					currentPlan[obj] = [0, 0, 0]

			if len(currentPlan) == self.config.taskNum + 1:
				feasiblePlan = True

		# execute
		# robot
		self.robot.robot_specific_reset(robotx, roboty, robotYaw, self.np_random)

		# objects
		colorList = [[0, 1, 0, 1], [1, 0, 0, 1], [0, 0, 1, 1], [1, 1, 0, 1]]
		for i, obj in enumerate(self.objList):  # for each object
			height = self.robot.entityZ if i in self.objInScene else -1
			if obj == 'cylinder':
				height = height - 0.1

			self._p.resetBasePositionAndOrientation(self.objUidList[i],
													[currentPlan[obj][0], currentPlan[obj][1], height],
													self._p.getQuaternionFromEuler([0, 0, currentPlan[obj][2]]))
			self.objPoseList.append(currentPlan[obj])

			# let's change color of the obj
			if self.config.domainRandomization:
				if self.np_random.rand() > 0.5:
					texChosen = self.np_random.choice(self.objTexList)
					self._p.changeVisualShape(self.objUidList[i], linkIndex=-1,
											  textureUniqueId=self.objTexDict[texChosen],
											  rgbaColor=list(self.rand_rgb(0.8)) + [1])
				else:
					self._p.changeVisualShape(self.objUidList[i], linkIndex=-1,
											  textureUniqueId=self.objTexDict['white_default.png'],
											  rgbaColor=list(self.np_random.rand(3, )) + [1])
			else:
				self._p.changeVisualShape(self.objUidList[i], linkIndex=-1,
										  rgbaColor=list(self.np_random.rand(3, )) + [1])
				self._p.changeVisualShape(self.objUidList[i], linkIndex=-1, rgbaColor=list(self.np_random.rand(3, )) + [
					1])  # let's always change obj color

		if self.config.domainRandomization:
			# change the texture of the outer arena
			self.changeWallTexture(self.arenaWallOuter)

			# change the texture of the inner arena
			self.changeWallTexture(self.arenaWallInner)

		self._p.stepSimulation()  # refresh the simulator. Needed for the ray test

	def randomXYYaw(self, rouLow, rouHigh, thetaLow, thetaHigh, yawLow, yawHigh):
		rou = self.np_random.uniform(low=rouLow, high=rouHigh)
		theta = self.np_random.uniform(low=thetaLow, high=thetaHigh)
		y = rou * np.sin(theta)
		x = rou * np.cos(theta)
		yaw = self.np_random.uniform(low=yawLow, high=yawHigh)
		return x, y, yaw

	def rand_rgb(self, val):
		"""
		:param val:
		:return: a list that contains r, g ,b values which satisfies following condition
				r + g + b >= 2 and r <= 1 and g <= 1 and b <= 1
		"""
		# generate 3 numbers in [0,1)
		rgb = self.np_random.rand(3)
		rgb = (1 - val) * rgb + val
		return rgb

	def changeWallTexture(self, wallID):
		texID = self.np_random.choice(self.wallTextureList)
		r, g, b = self.rand_rgb(2. / 3)  # generate random rgb values
		self._p.changeVisualShape(wallID, -1, textureUniqueId=texID,
								  rgbaColor=[r, g, b, 1])

	def envReset(self):
		if self.robot.robot_ids is None: # if it is the first run, load all models
			# load sound
			# if os is Linux, we already load the data in the parent process. See shmem_vec_env.py
			if len(self.audio.words) == 0:
				self.audio.loadData()

			self.loadArena()
			# load objects and put them far away from the arena
			for item in self.objList :
				self.objUidList.append(self.loadObj(os.path.join(self.config.mediaPath, 'objects', item + ".obj"), item))

			# load robot so that robot_ids is not None
			self.robot.load_model()
			self.robot._p=self._p
			self.robot.eyes=self.robot.parts["eyes"] # in the urdf, the link for the camera is called "eyes"
			self.objInScene=np.arange(len(self.objList))
			self.randomization()

		if self.config.ifReset and self.episodeCounter>0:
			self.randomization()

		ret=self.gen_obs()
		obs=ret[0]

		return obs

	def saveEpisodeImage(self, image):
		if self.config.episodeImgSaveInterval > 0 and self.episodeCounter % self.config.episodeImgSaveInterval == 0:
			# save the images
			imgSave = cv2.resize(image, (self.config.episodeImgSize[1], self.config.episodeImgSize[0]))
			if self.config.episodeImgSize[2] == 3:
				imgSave = cv2.cvtColor(imgSave, cv2.COLOR_RGB2BGR)
			fileName = str(self.givenSeed) + 'out' + str(self.envStepCounter) + '.jpg'
			cv2.imwrite(os.path.join(self.config.episodeImgSaveDir, fileName), imgSave)

	def isFeasiblePose(self, new_obj, currentPlan, objList, objInScene):
		"""
		Check if a randomized pose will physically intersect with the wall or objects
		new_obj is a tuple ('robot', [robotx, roboty, robotyaw]).
		currentPlan is a list which has the location of all the objects obtained from randomization()
		the order will be same as self.objList . list: [[x,y,yaw],[x,y,yaw],...].
		Meanwhile, get the distance information for calc_state()

		:return:
		True or False indicating if the pose is valid
		"""

		new_key, new_pose = new_obj
		tot_dist = self.config.robotExpandDistance + self.config.robotRadius

		# check against other objects
		for i,value in enumerate(currentPlan):
			if i in objInScene: # get information of objects in the scene
				# get distance between the robot and the object
				distance = np.linalg.norm([new_pose[0] - value[0], new_pose[1] - value[1]])
				threshold=self.config.robotRadius+self.config.objectsRadius[objList[i]] + self.config.objectsExpandDistance[objList[i]]
				if distance <= threshold:
					return False

		# check against wall
		if new_pose[0] - tot_dist <= -0.95 or new_pose[0] + tot_dist >= 0.95 or \
				new_pose[1] - tot_dist <= -0.95 or new_pose[1] + tot_dist >= 0.95:
			return False

		return True

	def reject_fn(self, new_obj, currentPlan):
		"""
		Check if an entity collides with the wall and all other entities at environment reset
		When a new episode starts, randomization() will be called.
		new_obj is the object going to be inserted. It is a tuple
		(obj_name in self.objList , purposed position=[x,y,yaw])
		currentPlan is a dict which contains the locations of all inserted entities. 'entity_name' includes 'robot'
		{'entity_name': [x,y,yaw]}

		The criterion of contact:
		We will model each entity as circles. The center of the circle is the center of the entity
		the radius of the circle is slightly bigger than the real radius of the entity
		e.g. The cube which has a dimension of 0.2*0.2 can be modeled as a circle with radius of 0.1+0.05,
		so that the robot will not really collide with the cube although the program says it does.
		Therefore, we are going to measure the distance between 2 entities.
		If the distance between 2 entities <= the sum of the radius of the two circles, the 2 entities are colliding

		:return:
		True or False indicating if the plan should be rejected
		"""

		new_key, new_pose = new_obj
		tot_dist = self.config.objectsExpandDistance[new_key] + self.config.objectsRadius[new_key]
		# check against wall
		if new_pose[0] - tot_dist <= -1 or new_pose[0] + tot_dist >= 1 or \
				new_pose[1] - tot_dist <= -1 or new_pose[1] + tot_dist >= 1:
			return True
		# check against other objects and the robot
		for key in currentPlan:
			# get distance between the robot and the object
			distance = np.linalg.norm([new_pose[0] - currentPlan[key][0], new_pose[1] - currentPlan[key][1]])
			if key == 'robot': # check against robot
				if distance <= self.config.robotRadius + self.config.robotExpandDistance + \
								self.config.objectsRadius[new_key] + self.config.objectsExpandDistance[new_key]+\
								self.config.placementExtension:
					return True
			else: # check against other objects
				if distance <= self.config.objectsRadius[key] + self.config.objectsRadius[new_key] + \
						self.config.objectsExpandDistance[key]+self.config.objectsExpandDistance[new_key] + \
						self.config.placementExtension:
					return True
		return False

	def get_positive_negative(self, inViewList, get_negative=True, generate_audio=True):
		positive_audio = None
		sound_positive=None
		sound_negative = None
		if sum(np.isinf(inViewList))!=self.config.taskNum-1: 
			# the agent sees nothing, no sound is given
			if generate_audio:
				sound_positive = np.zeros(shape=self.config.sound_dim)
			if get_negative:
				objIndx = self.np_random.randint(0, self.config.taskNum)
				sound_negative, _ = self.audio.genSoundFeat(objIndx=objIndx, featType='MFCC',
															rand_fn=self.np_random.randint)
			ground_truth = np.int32(self.config.taskNum)
		else:  # the agent sees an object in self.config.objList
			objIndx_positive = np.argmin(inViewList)  # choose the one with min distance
			ground_truth = np.int32(objIndx_positive)
			if generate_audio:
				sound_positive, positive_audio = self.audio.genSoundFeat(objIndx=objIndx_positive, featType='MFCC',
																		 rand_fn=self.np_random.randint)
			if get_negative:
				objIndx_negative = self.np_random.randint(0, self.config.taskNum)
				if objIndx_positive == objIndx_negative:
					sound_negative = np.zeros(shape=self.config.sound_dim)
				else:
					sound_negative, negative_audio = self.audio.genSoundFeat(objIndx=objIndx_negative, featType='MFCC',
																			 rand_fn=self.np_random.randint)
		return sound_positive, sound_negative, ground_truth, positive_audio

	def gen_obs(self):
		"""
		:return: a dict containing various type of observations
		"""
		image=self.robot.get_image()

		self.saveEpisodeImage(image)

		s = self.robot.calc_state(self.objPoseList, self.objList, self.objInScene)

		rayTest = self.robot.ray_test(self.objUidList)
		contactRays = [True if any((rayTest == Uid)[1:-1]) else False for Uid in self.objUidList]

		inViewList = [np.inf] * self.config.taskNum
		for i in range(len(self.objList)):
			if contactRays[i] and 0 <= s['dist'][i] <= 0.15:
				inViewList[i] = s['dist'][i]  # object not being seen has inf distance

		# sound_positive: the current sound heard by the agent
		# sound_negative: the sound that is not the current heard sound
		# sound_positive_ground_truth: the ground truth label for the sound heard by the agent
		# at training time, generate triplets from time to time
		get_negative = self.config.RLTrain and self.np_random.rand() > self.config.pretextModelUpdatePairProb
		sound_positive, sound_negative, sound_positive_ground_truth, _ = self.get_positive_negative(inViewList, get_negative)

		if self.envStepCounter==0:
			if self.config.RLTrain or self.config.render:
				self.goalObjIdx = self.np_random.randint(0, self.config.taskNum)
			else:
				idx = np.where(self.size_per_class_cumsum <= self.episodeCounter)[0]
				self.goalObjIdx = 0 if len(idx) == 0 else int(idx.max() + 1)
				self.goalObjIdx = np.clip(self.goalObjIdx, a_min=0, a_max=self.config.taskNum-1)

			self.goal_sound, self.goal_audio = self.audio.genSoundFeat(objIndx=self.goalObjIdx, featType='MFCC',
																	   rand_fn=self.np_random.randint)
			self.ground_truth = np.int32(self.goalObjIdx)
			if self.config.render or self.config.RLTrain == False:
				if self.goal_audio is not None and self.config.render:
					sd.play(self.goal_audio, self.audio.fs)
				print('Goal object is', self.goalObjIdx)

		obs = {
			'image': np.transpose(image, (2, 0, 1)),
			'goal_sound': self.goal_sound,
			'current_sound': sound_positive,
			'robot_pose': np.array([s['motor'][0], s['motor'][1]]),
			'ground_truth': self.ground_truth,
		}

		return obs, s, sound_positive, sound_negative, sound_positive_ground_truth, inViewList

	def step(self, action):
		action=np.array(action)
		if self.config.record and self.config.loadAction: # replace network output with other actions if needed
			action = self.episodeRecorder.loadedAction[self.envStepCounter, :]

		act=[]
		if self.config.RLRobotControl in ['rotPos']:
			if self.config.RLManualControl:
				k = p.getKeyboardEvents()
				dx=0.
				dtheta=0.0
				if p.B3G_UP_ARROW in k:
					dx = 0.05
				if p.B3G_DOWN_ARROW in k:
					dx = -0.05
				if p.B3G_LEFT_ARROW in k:
					dtheta = 0.15
				if p.B3G_RIGHT_ARROW in k:
					dtheta = -0.15
				act=[dx,dtheta]

			else:
				# action[0] transitional  velocity, action[1] rotational position
				deltaTrans = 0.05
				dTrans = float(np.clip(action[0], -1, +1)) * deltaTrans  # delta transitional velocity

				deltaRot = 0.15
				dRot = float(np.clip(action[1], -1, +1)) * deltaRot  # delta rotational position
				act = [dTrans, dRot]

		elif self.config.RLRobotControl in ['pointFollower', 'setPoseInc']:
			x = y = theta = None
			# action dimension should be 3
			max_xy = self.config.xyMax
			if self.config.RLManualControl:
				x = np.clip(float(input('desired x:')), -max_xy, max_xy)
				y = np.clip(float(input('desired y:')), -max_xy, max_xy)
				theta = np.clip(float(input('desired theta:')), -np.pi, np.pi)
				act = [x, y, theta]
			else:
				deltaX = 0.1
				dX = float(np.clip(action[0], -1, +1)) * deltaX

				deltaY = 0.1
				dY = float(np.clip(action[1], -1, +1)) * deltaY

				deltaTheta=0.5
				dth=float(np.clip(action[2], -1, +1)) * deltaTheta
				act = [dX, dY, dth]

		else:
			raise NotImplementedError

		real_action = self.robot.apply_action(act,
											  controlMethod=self.config.RLRobotControl,
											  currentPlan=self.objPoseList,
											  objList=self.objList,
											  objInScene=self.objInScene)
		self.scene.global_step()
		self.envStepCounter = self.envStepCounter + 1

		obs, s, sound_positive, sound_negative, sound_positive_ground_truth, inViewList = self.gen_obs()

		r =self.rewards(s) # calculate reward
		self.reward = sum(r)
		self.episodeReward = self.episodeReward + self.reward
		self.done = self.termination()

		infoDict={}

		if sound_negative is not None:
			infoDict['image'] = obs['image']
			infoDict['sound_positive'] = sound_positive.astype(np.float32)
			infoDict['sound_negative'] = sound_negative.astype(np.float32)
			infoDict['ground_truth'] = np.array([sound_positive_ground_truth], dtype=np.int32)

		if not self.config.RLTrain:
			if not np.isinf(inViewList[self.goalObjIdx]):
				self.goal_area_count=self.goal_area_count+1
			if self.done:
				infoDict['goal_area_count']=self.goal_area_count
				print('goal area count-------------------------', self.goal_area_count)
				self.goal_area_count = 0

		if self.config.record:
			self.episodeRecorder.wheelVelList.append(real_action)
			self.episodeRecorder.actionList.append(list(action))

		return obs, self.reward, self.done, infoDict # reset will be called if done

	def rewards(self, *args):
		return [0]

	def termination(self):
		done=False
		if self.envStepCounter >= self.maxSteps:
			done=True

		return done


class Recorder(object):
	def __init__(self):
		self.actionList = []  # recorded in turtlebot_env.step
		self.wheelVelList = []  # recorded in apply_action()
		self.orientationList = []  # recorded in calc_state()
		self.positionList = []  # recorded in calc_state()
		self.predictExi = []
		self.predictInSight = []
		self.episodeInitNum=0
		self.saveTo = None
		self.loadFrom = None
		self.loadedAction = None

		#################################
		self.reward = []  # records Episode Reward
		self.final_distance = []  # surface distance
		self.object = []  # The target object in current episode: 0->cube, 1->sphere, 2->cone, 3->cylinder
		self.distanceInArea=[]
		#################################

	def saveEpisode(self,episodeCounter):
		savePath = os.path.join(self.saveTo, 'ep' + str(self.episodeInitNum + episodeCounter))
		if not os.path.exists(savePath):
			os.makedirs(savePath)

		# episode info
		reward = pd.DataFrame({'target object': np.array(self.object), 'reward': np.array(self.reward),
							   'final dist to goal': np.array(self.final_distance),
							   'distanceInArea':np.array(self.distanceInArea)})

		reward.to_csv(os.path.join(savePath, 'results.csv'), index=False)
		print("csv written")
		self.clear()

	def loadActions(self):
		self.loadedAction = pd.read_csv(self.loadFrom)
		self.loadedAction = self.loadedAction.values
		print("Reading actions from", self.loadFrom)

	def clear(self):
		self.actionList = []  # recorded in turtlebot_env.step. [trans,rot]
		self.wheelVelList = []  # recorded in apply_action(). [left,right]
		self.orientationList = []  # recorded in calc_state() [angle]
		self.positionList = []  # recorded in calc_state() [x,y]
		self.predictExi = []
		self.predictInSight = []

		self.reward = []  # records Episode Reward
		self.final_distance = []  # surface distance
		self.object = []
		self.distanceInArea = []
