import os
import sys
from sys import platform
if platform == "linux" or platform == "linux2":
	import tty
	import termios
import numpy as np
import cv2
import sounddevice as sd
import gym
from ai2thor.controller import Controller

from ..RSI2.RL_env_RSI2 import RLEnvRSI2, Task


class RLEnvRSI1(RLEnvRSI2):

	def __init__(self):
		RLEnvRSI2.__init__(self)

		# update observation space
		d = {
			'image': self.observation_space['image'],
			'occupancy': self.observation_space['occupancy'],
			'goal_sound': self.observation_space['goal_sound'],
			# the ground_truth is the goal sound label
			'soundLabel': gym.spaces.Box(low=-np.inf, high=np.inf, shape=(self.config.taskNum,), dtype=np.int32),
			# one-hot encoding indicating which objects are being seen
			'inSight': gym.spaces.Box(low=-np.inf, high=np.inf, shape=(self.config.taskNum,), dtype=np.int32),
			# whether the goal object is being seen
			'exi': gym.spaces.Box(low=-np.inf, high=np.inf, shape=(1,), dtype=np.int32),
		}
		self.observation_space = gym.spaces.Dict(d)

		# for calculating the rewards
		self.potential = {'disPotential': 0., 'angPotential': 0.}

	def reset(self):
		# choose the task
		# decide the types of rooms and the floor_plan
		self.taskID = self.np_random.randint(len(self.taskList))
		self.task = self.taskList[self.taskID]
		self.floor_plan = self.np_random.choice(self.config.allScene[self.task.loc])

		if self.controller is None:  # if it is the first round

			visibleDist = 100.0 # for inSight and exi. For object interaction, we use RLVisibleDistance

			self.controller = Controller(agentMode="default", visibilityDistance=visibleDist,
										 scene='FloorPlan' + str(self.floor_plan), gridSize=self.config.gridSize[self.floor_plan],
										 snapToGrid=self.config.snapToGrid, continuous=True,
										 rotateStepDegrees=self.config.rotateStepDegrees, renderDepthImage=False,
										 renderInstanceSegmentation=False, width=self.config.img_dim[2],
										 height=self.config.img_dim[1], fieldOfView=self.config.fieldOfView)
		else:
			# reset variables
			self.episodeCounter = self.episodeCounter + 1
			self.done = 0
			self.reward = 0
			self.terminated = False
			self.envStepCounter = 0  # if max time step of this env is 150, then envStepCounter will be 150 before reset
			self.episodeReward = 0.0

			self.controller.reset('FloorPlan' + str(self.floor_plan))
			self.controller.step('ResetObjectFilter')
			self.objMeta = {}
			self.visibility = {}
			self.potential = {'disPotential': 0., 'angPotential': 0.}  # potential

		if self.floor_plan not in self.reachablePositions:
			pos = self.controller.step(action="GetReachablePositions").metadata["actionReturn"]
			self.reachablePositions[self.floor_plan] = []
			self.robotY[self.floor_plan] = pos[0]['y']
			for item in pos:
				self.reachablePositions[self.floor_plan].append((item['x'], item['z']))

			l = np.array(self.reachablePositions[self.floor_plan])
			self.min_xz[self.floor_plan] = np.min(l, axis=0)
			max_xz = np.max(l, axis=0)
			self.get_occupancy_grid(self.config.gridSize[self.floor_plan], max_xz)
		self.setupTask()

		ret = self.gen_obs()
		self.calc_potential()  # get initial potential
		return ret[0]

	def special_action(self, action_str):

		if action_str in ["ToggleObjectOn", "ToggleObjectOff"]:
			obj_in_view=None
			for k in self.visibility:
				if k!='Pillow':
					if self.visibility[k]:
						obj_in_view=k
			if obj_in_view is not None:
				distance = self.get_distance(obj_in_view)
				if distance <= self.config.RLVisibilityDistance:
					e = self.controller.step(action=action_str, objectId=self.objMeta[obj_in_view]["objectId"])
		elif action_str in ["PickupObject"]:
			# we perform forced pick up as long as the object is within a certain distance to the robot
			distance = self.get_distance('Pillow')
			if distance <= self.config.RLVisibilityDistance:
				e=self.controller.step(
					action="PickupObject",
					objectId=self.objMeta['Pillow']["objectId"],
					forceAction=False,
					manualInteract=False
				)
		else:
			raise NotImplementedError

	def gen_obs(self):
		"""
		:return: a dict containing various type of observations
		"""

		# update object metadata
		self.updateObjMeta(list(self.objMeta.keys()))
		self.checkVisible()
		self.agentMeta=self.controller.last_event.metadata["agent"]
		rgb_image=self.controller.last_event.frame
		self.saveEpisodeImage(rgb_image)
		image=rgb_image

		self.local_occupancy=self.get_local_occupancy_map(x=self.agentMeta['position']['x'],
											   z=self.agentMeta['position']['z'],
											   y=self.agentMeta['rotation']['y'])

		inSight = []
		for obj in self.config.allTasks[self.task.loc]:
			if obj in ['FloorLamp', 'Television']:
				l = [0] * len(self.config.allTasks[self.task.loc][obj])
				if self.visibility[obj]:
					if self.objMeta[obj]["isToggled"]:
						l[0] = 1
					else:
						l[1] = 1
				inSight.extend(l)
			elif obj == 'Pillow':
				if self.visibility[obj]:
					inSight.append(1)
				else:
					inSight.append(0)
			else:
				raise NotImplementedError

		exi=0
		if self.visibility[self.task.obj]:
			exi=1

		if self.envStepCounter==0:
			if self.config.RLTrain or self.config.render:
				# select an audio according to the task
				self.goal_sound, self.goal_audio, self.transcription=self.audio.getAudioFromTask(self.np_random, self.task, Task)
			else:
				self.goal_sound, self.goal_audio, self.transcription=self.audio.getAudioFromTask(self.np_random, self.task, Task)

			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 intent is', self.task.loc+' '+self.task.obj+' '+self.task.act)

		else:
			self.goal_sound=np.ones_like(self.goal_sound)*np.inf

		soundLabel = np.zeros((self.config.taskNum,), dtype=np.int32)
		soundLabel[self.task2ID[self.task]]=1

		obs = {
			'image': np.transpose(image, (2, 0, 1)),
			'occupancy':np.transpose(np.expand_dims(self.local_occupancy, -1), (2,0,1)),
			'goal_sound': self.goal_sound,
			'soundLabel': soundLabel,
			'inSight': np.array(inSight),
			'exi': np.array(exi),
		}

		return obs, None, None

	def rewards(self, *args):
		distanceProgress = 0
		potential_old = self.potential
		self.calc_potential()  # update potential
		# distance
		distance_factor = self.config.distanceRewardFactor
		distanceProgress = distanceProgress + distance_factor * (
				self.potential['disPotential'] - potential_old['disPotential'])

		goal_reward=1 if self.checkTaskDone() else 0
		return [distanceProgress, goal_reward]

	def get_distance(self, obj):
		robot_x = self.agentMeta['position']['x']
		robot_z = self.agentMeta['position']['z']

		target_x = self.objMeta[obj]['position']['x']
		target_z = self.objMeta[obj]['position']['z']

		distance = np.linalg.norm([robot_x - target_x, robot_z - target_z])
		return distance


	def calc_potential(self):
		# get distance
		distance=self.get_distance(self.task.obj)

		self.potential = {'disPotential': 0., 'angPotential': 0.}
		# the potential has two parts: a part for distance and a part for angle
		self.potential['disPotential']=-abs(distance)
