"""
Environment for 2D car driving.
You can customize this script in a way you want.

View more on [莫烦Python] : https://morvanzhou.github.io/tutorials/


Requirement:
pyglet >= 1.2.4
numpy >= 1.12.1
"""
import numpy as np
import pyglet
import random as rand
import maps
from collections import deque

pyglet.clock.set_fps_limit(10000)


class CarEnv(object):
	n_sensor = 13
	action_dim = 1
	actions = [ -1, 0, 1 ]
	n_actions = len(actions)
	state_dim = n_sensor
	sensor_max = 120.
	
	speed = 120.
	dt = 0.2

	def __init__(self, discrete_action=False, step_back=False, map_set=0):
		self.is_discrete_action = discrete_action
		self.step_back = step_back
		self.step_count = 0

		self.viewer = None
		self.map_ = maps.obstacles(map_set=map_set)
		self.viewer_xy = self.map_.choose_map_size()
		self.start_point = self.map_.choose_start_point()

		if discrete_action:
			pass
			#self.actions = [-1.2, -1, -0.5, 0, 0.5, 1, 1.2]
		else:
			self.action_bound = [-2, 2]

		self.terminal = False
		# node1 (x, y, r, w, l),
		self.car_info = np.array([0, 0, 0, 20, 20], dtype=np.float64)   # car coordination

		self.obstacle_coords = self.map_.choose_map()

		self.n_obsacles = len(self.obstacle_coords)

		self.sensor_info = self.sensor_max + np.zeros((self.n_sensor, 3))  # n sensors, (distance, end_x, end_y)

		self.set_fps(fps=200)

		self.start = True

		self.die_position=deque(maxlen=8)
		self.current_position=deque(maxlen=8)


	def step(self, action):
		
		if self.is_discrete_action:
			action = self.actions[action]
		else:
			action = np.clip(action, *self.action_bound)[0]

		
		
		if abs(action) != 2 :
			
			self.car_info[2] += action * np.pi/30  # max r = 6 degree	
			self.car_info[:2] = self.car_info[:2] + self.speed * self.dt * np.array([np.cos(self.car_info[2]), np.sin(self.car_info[2])])

			self.step_count += 1
			self.current_position.append(np.copy(self.car_info[:3]))
		else :
			self.car_info[2] += action * np.pi/30 * 1.2


		self._update_sensor()
		s = self._get_state()
		r = -1 if self.terminal else 1

		return s, r, self.terminal

	def reset(self):
		self.terminal = False
		out = not 0 < self.car_info[0] < self.viewer_xy[0] or not 0 < self.car_info[1] < self.viewer_xy[1]
		if self.start or out or np.random.random() < 0.05 or not self.step_back: 
			self.car_info[:3] = np.array([*self.start_point[rand.randrange(0, len(self.start_point))]])
			self.start=False
		else : 
			if self.step_back:
				if self.step_count < 8: 
					pass
				else:
					for i in range(len(self.current_position)): 
						self.die_position.append(self.current_position[i])
				self.car_info[:3] = self.die_position[0]
			else:
				self.car_info[:2] = self.car_info[:2] + 3 * self.speed * self.dt * np.array([np.cos(self.car_info[2]), np.sin(self.car_info[2])])
		self._update_sensor()
		self.step_count = 0
		return self._get_state()

	def render(self):
		if self.viewer is None:
			self.viewer = Viewer(*self.viewer_xy, self.car_info, self.sensor_info, self.obstacle_coords)
		self.viewer.render()

	def sample_action(self):
		if self.is_discrete_action:
			a = np.random.choice(list(range(3)))
		else:
			a = np.random.uniform(*self.action_bound, size=self.action_dim)
		return a

	def set_fps(self, fps=30):
		pyglet.clock.set_fps_limit(fps)

	def _get_state(self):
		s = self.sensor_info[:, 0].flatten()/self.sensor_max
		return s

	def check_dead(self):
		out = False
		for i in range(len(self.obstacle_coords)):
			for j in range(len(self.obstacle_coords[i])):
				pass

	def _update_sensor(self):
		cx, cy, rotation = self.car_info[:3]

		n_sensors = len(self.sensor_info)
		sensor_theta = np.linspace(-np.pi / 2, np.pi / 2, n_sensors)
		xs = cx + (np.zeros((n_sensors, ))+self.sensor_max) * np.cos(sensor_theta)
		ys = cy + (np.zeros((n_sensors, ))+self.sensor_max) * np.sin(sensor_theta)
		xys = np.array([[x, y] for x, y in zip(xs, ys)])	# shape (5 sensors, 2)

		# sensors
		tmp_x = xys[:, 0] - cx
		tmp_y = xys[:, 1] - cy
		# apply rotation
		rotated_x = tmp_x * np.cos(rotation) - tmp_y * np.sin(rotation)
		rotated_y = tmp_x * np.sin(rotation) + tmp_y * np.cos(rotation)
		# rotated x y
		self.sensor_info[:, -2:] = np.vstack([rotated_x+cx, rotated_y+cy]).T

		q = np.array([cx, cy])
		for si in range(len(self.sensor_info)):
			s = self.sensor_info[si, -2:] - q
			possible_sensor_distance = [self.sensor_max]
			possible_intersections = [self.sensor_info[si, -2:]]

			# obstacle collision
			for i in range(len(self.obstacle_coords)):
				for oi in range(len(self.obstacle_coords[i])):
					p = self.obstacle_coords[i][oi]
					r = self.obstacle_coords[i][(oi + 1) % len(self.obstacle_coords[i])] - self.obstacle_coords[i][oi]
					if np.cross(r, s) != 0:  # may collision
						t = np.cross((q - p), s) / np.cross(r, s)
						u = np.cross((q - p), r) / np.cross(r, s)
						if 0 <= t <= 1 and 0 <= u <= 1:
							intersection = q + u * s
							possible_intersections.append(intersection)
							possible_sensor_distance.append(np.linalg.norm(u*s))

			# window collision
			win_coord = np.array([
				[0, 0],
				[self.viewer_xy[0], 0],
				[*self.viewer_xy],
				[0, self.viewer_xy[1]],
				[0, 0],
			])
			for oi in range(4):
				p = win_coord[oi]
				r = win_coord[(oi + 1) % len(win_coord)] - win_coord[oi]
				if np.cross(r, s) != 0:  # may collision
					t = np.cross((q - p), s) / np.cross(r, s)
					u = np.cross((q - p), r) / np.cross(r, s)
					if 0 <= t <= 1 and 0 <= u <= 1:
						intersection = p + t * r
						possible_intersections.append(intersection)
						possible_sensor_distance.append(np.linalg.norm(intersection - q))

			distance = np.min(possible_sensor_distance)
			distance_index = np.argmin(possible_sensor_distance)
			self.sensor_info[si, 0] = distance
			self.sensor_info[si, -2:] = possible_intersections[distance_index]

			out = not 0 < self.car_info[0] < self.viewer_xy[0] or not 0 < self.car_info[1] < self.viewer_xy[1]
			if distance <= self.car_info[-1]/1.5 or out:
				self.terminal = True


class Viewer(pyglet.window.Window):
	color = {
		'background': [1]*3 + [1]
	}
	fps_display = pyglet.clock.ClockDisplay()
	bar_thc = 5

	def __init__(self, width, height, car_info, sensor_info, obstacle_coords):
		super(Viewer, self).__init__(width, height, resizable=False, caption='2D car', vsync=False)  # vsync=False to not use the monitor FPS
		self.set_location(x=80, y=10)
		pyglet.gl.glClearColor(*self.color['background'])

		self.car_info = car_info
		self.sensor_info = sensor_info

		self.batch = pyglet.graphics.Batch()
		background = pyglet.graphics.OrderedGroup(0)
		foreground = pyglet.graphics.OrderedGroup(1)

		self.sensors = []
		line_coord = [0, 0] * 2
		c = (73, 73, 73) * 2
		for i in range(len(self.sensor_info)):
			self.sensors.append(self.batch.add(2, pyglet.gl.GL_LINES, foreground, ('v2f', line_coord), ('c3B', c)))

		car_box = [0, 0] * 4
		c = (249, 86, 86) * 4
		self.car = self.batch.add(4, pyglet.gl.GL_QUADS, foreground, ('v2f', car_box), ('c3B', c))

		c = (134, 181, 244) * 4
		self.obstacle_coords = obstacle_coords
		for i in range(len(obstacle_coords)):
			self.obstacle = self.batch.add(4, pyglet.gl.GL_QUADS, background, ('v2f', obstacle_coords[i].flatten()), ('c3B', c))

	def render(self):
		pyglet.clock.tick()
		self._update()
		self.switch_to()
		self.dispatch_events()
		self.dispatch_event('on_draw')
		self.flip()

	def on_draw(self):
		self.clear()
		self.batch.draw()
		# self.fps_display.draw()

	def _update(self):
		cx, cy, r, w, l = self.car_info

		# sensors
		for i, sensor in enumerate(self.sensors):
			sensor.vertices = [cx, cy, *self.sensor_info[i, -2:]]

		# car
		xys = [
			[cx + l / 2, cy + w / 2],
			[cx - l / 2, cy + w / 2],
			[cx - l / 2, cy - w / 2],
			[cx + l / 2, cy - w / 2],
		]
		r_xys = []
		for x, y in xys:
			tempX = x - cx
			tempY = y - cy
			# apply rotation
			rotatedX = tempX * np.cos(r) - tempY * np.sin(r)
			rotatedY = tempX * np.sin(r) + tempY * np.cos(r)
			# rotated x y
			x = rotatedX + cx
			y = rotatedY + cy
			r_xys += [x, y]
		self.car.vertices = r_xys


if __name__ == '__main__':
	np.random.seed(1)
	env = CarEnv()
	env.set_fps(30)
	for ep in range(20):
		s = env.reset()
		# for t in range(100):
		while True:
			env.render()
			s, r, done = env.step(env.sample_action())
			if done:
				break
