import gym
import numpy as np
import pybullet as p
import cv2
import gym, gym.spaces, gym.utils, gym.utils.seeding
from pybullet_utils import bullet_client
import os
try:
	if os.environ["PYBULLET_EGL"]:
		import pkgutil
except:
	pass


class BaseEnv(gym.Env):
	"""
	Base class for Bullet physics simulation environments in a Scene.
	These environments create single-player scenes and behave like normal Gym environments, if
	you don't use multiplayer.
	"""

	def __init__(self, robot, config, render, action_space, observation_space):
		# Pybullet related
		self.scene = None
		self.physicsClientId = -1 # at the first run, we do not own physics client
		self.ownsPhysicsClient = False
		self._p=None
		self.config=config
		self.timeStep = 1. / 240.

		# setup GUI camera
		self.debugCam_dist=config.debugCam_dist
		self.debugCam_yaw=config.debugCam_yaw
		self.debugCam_pitch=config.debugCam_pitch

		# robot related
		self.isRender = render
		self.robot = robot
		self.action_space = action_space
		self.observation_space = observation_space

		#episode related
		self.episodeCounter=-1
		self.envStepCounter=0
		self.done = 0
		self.reward = 0
		self.episodeReward = 0.0
		self.terminated = False

		self.np_random=None
		self.givenSeed=None

		# Debug
		self.logID = None

	def create_single_player_scene(self, p):
		raise NotImplementedError

	def seed(self, seed=None):
		# use a random seed if seed is None
		self.np_random, seed = gym.utils.seeding.np_random(seed)
		self.givenSeed=seed

		print("Created a Kuka Env with seed", seed)
		return [seed]

	def reset(self):
		# After gym.make, we call env.reset(). This function is then called.
		# this function starts pybullet client and create_single_player_scene
		# if it is the first run, setup Pybullet client and set GUI camera
		if self.physicsClientId < 0:
			self.ownsPhysicsClient = True

			if self.isRender:
				self._p = bullet_client.BulletClient(connection_mode=p.GUI)
			else:
				self._p = bullet_client.BulletClient()

			self._p.resetSimulation()
			self._p.setPhysicsEngineParameter(deterministicOverlappingPairs=1)

			# optionally enable EGL for faster headless rendering
			try:
				if os.environ["PYBULLET_EGL"]:
					con_mode = self._p.getConnectionInfo()['connectionMethod']
					if con_mode == self._p.DIRECT:
						egl = pkgutil.get_loader('eglRenderer')
						if (egl):
							self._p.loadPlugin(egl.get_filename(), "_eglRendererPlugin")
						else:
							self._p.loadPlugin("eglRendererPlugin")
			except:
				pass

			self.physicsClientId = self._p._client
			self._p.configureDebugVisualizer(p.COV_ENABLE_GUI,0)
			self._p.resetDebugVisualizerCamera(cameraDistance=self.debugCam_dist, cameraYaw=self.debugCam_yaw,
												cameraPitch=self.debugCam_pitch, cameraTargetPosition=[0, 0, 0])
			assert self._p.isNumpyEnabled() == 1

		# if it is the first run, build scene and setup simulation physics
		if self.scene is None:
			# this function will call episode_restart()->clean_everything() in
			# class World in scene_bases.py (self.cpp_world)
			# set gravity, set physics engine parameters, etc.
			self.scene = self.create_single_player_scene(self._p)

		# if it is not the first run
		if self.ownsPhysicsClient:
			self.scene.episode_restart()
		# reset counters
		self.episodeCounter = self.episodeCounter + 1
		self.done = False
		self.reward = 0
		self.terminated = False
		# if max time step of this env is 150, then envStepCounter will be 150 before reset
		self.envStepCounter = 0
		self.episodeReward = 0.0

		obs=self.envReset()
		return obs

	def envReset(self):
		raise NotImplementedError

	def render(self, mode='human'):
		# no need to implement this function
		pass

	def close(self):
		if self.ownsPhysicsClient:
			if self.physicsClientId >= 0:
				self._p.disconnect()
		self.physicsClientId = -1

	def step(self, action):
		raise NotImplementedError
