# First-Visit MC Policy Improvement
# without Exploring Starts (Epsilon-Greedy)
# Author: Yiheng Zhu
# Date: 14th/Dec/2017

import numpy as np
import Gridworld

# hyperparameters
GAMMA = 0.9 # discount rate
NUM_EPISODES = 100 # number of episodes
ACTIONS = ('U', 'D', 'L', 'R') # all possible actions

# find the key and value of the maximum value in a dictionary
def max_dict(dict_tmp):
	max_key = None
	max_val = float('-inf')
	for key, val in dict_tmp.items():
		if val > max_val:
			max_val = val
			max_key = key
	return max_key, max_val

def random_action(a, eps=0.1):
	p = np.random.random()
	if p < (1 - eps):
		return a
	else:
		return np.random.choice(ACTIONS)

def play_one_episode(gridworld, policy):
	# fixed starting state
	s = (2, 0)
	gridworld.set_state(s)
	a = random_action(policy[s])

	states_actions_rewards = [(s, a, 0)]
	while True:
		r = gridworld.step(a)
		s = gridworld.get_state()
		if gridworld.is_terminal():
			states_actions_rewards.append((s, None, r))
			break
		else:
			a = random_action(policy[s])
			states_actions_rewards.append((s, a, r))

	# calculate the return
	G = 0 # initialize the return, G, as 0
	states_actions_returns = [] # an empty list to store s, a and G
	first_step = True # first step is already in the list [(s, a, 0)]
	# calculate and store returns from terminal state to start state
	for s, a, r in reversed(states_actions_rewards):
		if first_step:
			# do nothing for the first step
			first_step = False
		else:
			# append state and return to the list
			states_actions_returns.append((s, a, G))
		G = r + GAMMA*G # update return
	# adjust the order of the list
	# (terminal -> start) -> (start -> terminal)
	states_actions_returns.reverse()
	return states_actions_returns

if __name__ == '__main__':
	# 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 |
	#    ---------------------
	gridworld = Gridworld.default_gridworld()

	# all traversable states
	states = gridworld.all_states()

	# draw the reward function
	print("Reward function:")
	Gridworld.draw_value_function(gridworld.rewards, gridworld)

	# initialize a random policy as a dictionary
	# for all traversable non-terminal states
	# policy = {}
	# for s in gridworld.actions.keys():
	# 	policy[s] = np.random.choice(ACTIONS)

	# inital policy
	# use the same initial condition
	# to compare with value iteration
	# -----------------
	# | D | L | D | T |
	# -----------------
	# | D | X | R | T |
	# -----------------
	# | U | L | L | U |
	# -----------------
	# X -- untraversable
	# T -- terminal
	policy = {
		(0, 0): 'D',
		(0, 1): 'L',
		(0, 2): 'D',
		(1, 0): 'D',
		(1, 2): 'R',
		(2, 0): 'U',
		(2, 1): 'L',
		(2, 2): 'L',
		(2, 3): 'U',
	}

	# initialize Q(s, a) and returns
	Q = {} # value function
	returns = {} # returns given state and action
	for s in states:
		if s in gridworld.actions.keys():
			Q[s] = {}
			for a in ACTIONS:
				Q[s][a] = 0
				returns[(s, a)] = []
		else:
			# non-traversable states and terminal states
			pass

	# repeat until convergence
	deltas = []
	for t in range(NUM_EPISODES):
		if t % 100 == 0:
			print(t)

		# generate an episode using policy pi
		delta = 0 # biggest change
		states_actions_returns = play_one_episode(gridworld, policy)
		visited_state_action_pairs = set()
		for s, a, G in states_actions_returns:
			# first-visit MC policy evaluation
			sa = (s, a)
			if sa not in visited_state_action_pairs:
				old_q = Q[s][a]
				returns[sa].append(G)
				Q[s][a] = np.mean(returns[sa])
				delta = max(delta, np.abs(old_q - Q[s][a]))
				visited_state_action_pairs.add(sa)
		deltas.append(delta)

		# update policy
		for s in policy.keys():
			policy[s] = max_dict(Q[s])[0] # only need the key

	# draw the final policy
	print("Final policy:")
	Gridworld.draw_policy(policy, gridworld)

	# calculate and draw the final value function
	V = {}
	for s, Qs in Q.items():
		V[s] = max_dict(Q[s])[1]
	print("Final value function:")
	Gridworld.draw_value_function(V, gridworld)
