# TD(0)
# Policy evaluation
# Uncertain state transition

import numpy as np
import matplotlib.pyplot as plt
from grid_world import standard_grid, negative_grid
from dp_ipe_dst_dp import print_values, print_policy

SMALL_ENOUGH = 1e-4
GAMMA = 0.9
ALPHA = 0.1

def affected_action(grid, s, a, eps=0.1):
	p = np.random.random()
	if p < (1 - eps) or len(grid.actions[s]) == 1:
		return a
	else:
		tmp = list(grid.actions[s])
		tmp.remove(a)
		return np.random.choice(tmp)

def play_one_episode(grid, policy):
	s = (0, 0)
	grid.set_state(s)
	states_and_rewards = [(s, 0)]
	while not grid.game_over():
		a_desired = policy[s]
		a = affected_action(grid, s, a_desired)
		r = grid.move(a)
		s = grid.current_state()
		states_and_rewards.append((s, r))
	return states_and_rewards

if __name__ == '__main__':
	grid = standard_grid()

	print("Rewards:")
	print_values(grid.rewards, grid)

	# ----------------
	# |R |R |R |R |f+|
	# ----------------
	# |D |xx|R |R |f-|
	# ----------------
	# |R |R |U |xx|U |
	# ----------------
	# |U |xx|U |R |U |
	# ----------------
	policy = {
		(0, 0): ('R'),
		(1, 0): ('R'),
		(2, 0): ('R'),
		(3, 0): ('R'),
		(0, 1): ('D'),
		(2, 1): ('R'),
		(3, 1): ('R'),
		(0, 2): ('R'),
		(1, 2): ('R'),
		(2, 2): ('U'),
		(4, 2): ('U'),
		(0, 3): ('U'),
		(2, 3): ('U'),
		(3, 3): ('R'),
		(4, 3): ('U'),
	}

	V = {}
	states = grid.all_states()
	for s in states:
		V[s] = 0

	for i in range(1000):
		states_and_rewards = play_one_episode(grid, policy)
		for t in range(len(states_and_rewards) - 1):
			s, _ = states_and_rewards[t]
			s2, r = states_and_rewards[t+1]
			V[s] = V[s] + ALPHA*(r + GAMMA*V[s2] - V[s])

	print("Values:")
	print_values(V, grid)
	print("Policy:")
	print_policy(policy, grid)
