# Gridworld
# Author: Yiheng Zhu
# Date: 11th/Nov/2017

import numpy as np

# gridworld environment
class Gridworld:
	def __init__(self, height, width, start):
		self.height = height # number of rows
		self.width = width # number of columns
		self.y = start[0] # start row
		self.x = start[1] # start column

	# set rewards and possible actions for all states in gridworld
	def init(self, rewards, actions):
		self.rewards = rewards # dictionary of rewards r (row, column)
		self.actions = actions # dictionary of actions A (row, column)

	# set current state
	def set_state(self, s):
		self.y = s[0]
		self.x = s[1]

	# get current state
	def get_state(self):
		return (self.y, self.x)

	# test if the state is terminal
	def is_terminal(self):
		return self.get_state() not in self.actions

	# move one step and get reward
	def step(self, action):
		# sanity check
		if action in self.actions[(self.y, self.x)]:
			if action == 'U':
				self.y -= 1
			elif action == 'D':
				self.y += 1
			elif action == 'L':
				self.x -= 1
			elif action == 'R':
				self.x += 1
		# return a reward if there is one (0 as default value)
		return self.rewards.get((self.y, self.x), 0)

	# return all states
	def all_states(self):
		return set(list(self.actions.keys()) + list(self.rewards.keys()))

# create a default gridworld
# oo -- traversable state
# xx -- untraversable state
# ss -- start state for agent
# +1 -- winning state with reward +1
# -1 -- losing state with reward -1
#      x0   x1   x2   x3
#    ---------------------
# y0 | oo | oo | oo | +1 |
#    ---------------------
# y1 | oo | xx | oo | -1 |
#    ---------------------
# y2 | ss | oo | oo | oo |
#    ---------------------
def default_gridworld():
	# init parameters
	height = 3
	width = 4
	start_y = 2
	start_x = 0
	start_state = (start_y, start_x)

	# dictionaries of rewards and actions
	# (column, row)
	# (y, x)
	rewards = {
		(0, 3): 1,
		(1, 3): -1,
		}
	actions = {
		(0, 0): ('D', 'R'),
		(0, 1): ('L', 'R'),
		(0, 2): ('L', 'D', 'R'),
		(1, 0): ('U', 'D'),
		(1, 2): ('U', 'D', 'R'),
		(2, 0): ('U', 'R'),
		(2, 1): ('L', 'R'),
		(2, 2): ('L', 'R', 'U'),
		(2, 3): ('L', 'U'),
	}

	# create a default gridworld
	default_gridworld = Gridworld(height, width, start_state)
	
	# initialize rewards and actions for default gridworld
	default_gridworld.init(rewards, actions)

	# return default gridworld
	return default_gridworld

# create a fast gridworld
# penalize movement to encourage fast termination
def fast_gridworld(step_cost=-0.1):
	# fast gridworld = default gridworld + step costs for non-terminal states
	fast_gridworld = default_gridworld()
	fast_gridworld.rewards.update({
		(0, 0): step_cost,
		(0, 1): step_cost,
		(0, 2): step_cost,
		(1, 0): step_cost,
		(1, 2): step_cost,
		(2, 0): step_cost,
		(2, 1): step_cost,
		(2, 2): step_cost,
		(2, 3): step_cost,
	})
	
	# return fast gridworld
	return fast_gridworld

# draw value function in command line
def draw_value_function(values, gridworld):
	for y in range(gridworld.height):
		print("------------------------")
		for x in range(gridworld.width):
			value = values.get((y, x), 0)
			# "-" sign takes up an extra space
			if value >= 0:
				print(" %.2f|" % value, end="")
			else:
				print("%.2f|" % value, end="")
		print("")
	print("------------------------")
	print("")

# draw policy in command line
def draw_policy(policy, gridworld):
	for y in range(gridworld.height):
		print("------------------------")
		for x in range(gridworld.width):
			action = policy.get((y, x), ' ')
			print("  %s  |" % action, end="")
		print("")
	print("------------------------")
	print("")