# Monte Carlo policy iteration
# Deterministic state transition
# First-visit
# Exploring starts

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

def play_one_episode(grid, policy):
	valid_states = list(grid.actions.keys())
	start_state_index = np.random.choice(len(valid_states))
	s = grid.set_state(valid_states[start_state_index])
	a = np.random.choice(list(policy[s]))

	states_actions_rewards = [(s, a, 0)]
	visited_states = set()

	while True:
		old_s = grid.current_state()
		r = grid.move(a)
		s = grid.current_state()

		if s in visited_states:
			states_actions_rewards.append((s, None, -100))
			break
		elif grid.game_over():
			states_actions_rewards.append((s, None, r))
			break
		else:
			a = policy[s]
			states_actions_rewards.append((s, a, r))
		visited_states.add(s)

	G = 0
	states_actions_returns = []
	first = True
	for s, a, r in reversed(states_actions_rewards):
		if first:
			first = False
		else:
			states_actions_returns.append((s, a, G))
		G = r + GAMMA * G
	states_actions_returns.reverse()
	return states_actions_returns

def max_dict(d):
	max_key = None
	max_value = float('-inf')
	for key, value in d.items():
		if value > max_value:
			max_value = value
			max_key = key
	return max_key, max_value

if __name__ == '__main__':
	grid = negative_grid(step_cost=-0.1)

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

	policy = {}
	for s in grid.actions.keys():
		policy[s] = np.random.choice(list(grid.actions[s]))

	Q = {}
	returns = {}
	states = grid.all_states()
	for s in grid.actions.keys():
		Q[s] = {}
		for a in grid.actions[s]:
			Q[s][a] = 0
			returns[(s, a)] = []

	deltas = []
	for i in range(2000):
		if i % 100 == 0:
			print(i)

		biggest_change = 0
		states_actions_returns = play_one_episode(grid, policy)
		visited_state_action_pairs = set()

		for s, a, G in states_actions_returns:
			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])
				biggest_change = max(biggest_change, np.abs(old_q - Q[s][a]))
				visited_state_action_pairs.add(sa)
		deltas.append(biggest_change)

		for s in policy.keys():
			policy[s] = max_dict(Q[s])[0]

	# plt.plot(deltas)
	# plt.show()

	print("Final policy:")
	print_policy(policy, grid)

	V = {}
	for s, Qs in Q.items():
		V[s] = max_dict(Q[s])[1]

	print("Final values:")
	print_values(V, grid)
