import os
import numpy as np
import cv2
import sounddevice as sd
import pandas as pd
import gym
from ..RSI2.RL_env_RSI2 import RLEnvRSI2
from ..env_bases import BaseEnv
import pybullet as p
from ..scene_abstract import SingleRobotEmptyScene


class RLEnvRSI3(RLEnvRSI2):
	def __init__(self):
		RLEnvRSI2.__init__(self)
		# 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),
			'goal_sound_label': gym.spaces.Box(low=0, high=4, shape=(1,), dtype=np.int32),
			'goal_sound_feat': gym.spaces.Box(low=-np.inf, high=np.inf, shape=(self.config.representationDim,),
															  dtype=np.float32),
			'image_feat': gym.spaces.Box(low=-np.inf, high=np.inf, shape=(self.config.representationDim,),
																			dtype=np.float32),
			'exi': gym.spaces.Box(low=-np.inf, high=np.inf, shape=(1,),
										 dtype=np.float32),
		}

		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)


	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, sound_negative, sound_positive_ground_truth, _ = \
			self.get_positive_negative(inViewList, get_negative=False, generate_audio=True)

		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]]),
			'goal_sound_label': self.ground_truth,
			'goal_sound_feat': np.zeros((self.config.representationDim,)),
			'image_feat': np.zeros((self.config.representationDim,)),
			'exi': np.zeros((1,))
		}

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