import gym
import numpy as np
from itertools import product
import matplotlib.pyplot as plt


def print_policy(Q, env):
    """ This is a helper function to print a nice policy from the Q function"""
    moves = [u'←', u'↓',u'→', u'↑']
    if not hasattr(env, 'desc'):
        env = env.env
    dims = env.desc.shape
    policy = np.chararray(dims, unicode=True)
    policy[:] = ' '
    for s in range(len(Q)):
        idx = np.unravel_index(s, dims)
        policy[idx] = moves[np.argmax(Q[s])]
        if env.desc[idx] in ['H', 'G']:
            policy[idx] = u'·'
    print('\n'.join([''.join([u'{:2}'.format(item) for item in row]) 
        for row in policy]))


def plot_V(Q, env):
    """ This is a helper function to plot the state values from the Q function"""
    fig = plt.figure()
    if not hasattr(env, 'desc'):
        env = env.env
    dims = env.desc.shape
    V = np.zeros(dims)
    for s in range(len(Q)):
        idx = np.unravel_index(s, dims)
        V[idx] = np.max(Q[s])
        if env.desc[idx] in ['H', 'G']:
            V[idx] = 0.
    plt.imshow(V, origin='upper', 
               extent=[0,dims[0],0,dims[1]], vmin=.0, vmax=.6, 
               cmap=plt.cm.RdYlGn, interpolation='none')
    for x, y in product(range(dims[0]), range(dims[1])):
        plt.text(y+0.5, dims[0]-x-0.5, '{:.3f}'.format(V[x,y]),
                horizontalalignment='center', 
                verticalalignment='center')
    plt.xticks([])
    plt.yticks([])


def plot_Q(Q, env):
    """ This is a helper function to plot the Q function """
    from matplotlib import colors, patches
    fig = plt.figure()
    ax = fig.gca()

    if not hasattr(env, 'desc'):
        env = env.env
    dims = env.desc.shape

    up = np.array([[0, 1], [0.5, 0.5], [1,1]])
    down = np.array([[0, 0], [0.5, 0.5], [1,0]])
    left = np.array([[0, 0], [0.5, 0.5], [0,1]])
    right = np.array([[1, 0], [0.5, 0.5], [1,1]])
    tri = [left, down, right, up]
    pos = [[0.2, 0.5], [0.5, 0.2], [0.8, 0.5], [0.5, 0.8]]
    
    cmap = plt.cm.RdYlGn
    norm = colors.Normalize(vmin=.0,vmax=.6)
    
    ax.imshow(np.zeros(dims), origin='upper', extent=[0,dims[0],0,dims[1]], vmin=.0, vmax=.6, cmap=cmap)
    ax.grid(which='major', color='black', linestyle='-', linewidth=2)

    for s in range(len(Q)):
        idx = np.unravel_index(s, dims)
        x, y = idx
        if env.desc[idx] in ['H', 'G']:
            ax.add_patch(patches.Rectangle((y, 3-x), 1, 1, color=cmap(.0)))
            plt.text(y+0.5, dims[0]-x-0.5, '{:.2f}'.format(.0),
                horizontalalignment='center', 
                verticalalignment='center')
            continue
        for a in range(len(tri)):
            ax.add_patch(patches.Polygon(tri[a] + np.array([y, 3-x]), color=cmap(Q[s][a])))
            plt.text(y+pos[a][0], dims[0]-1-x+pos[a][1], '{:.2f}'.format(Q[s][a]), 
                     horizontalalignment='center', verticalalignment='center',
                    fontsize=9, fontweight=('bold' if Q[s][a] == np.max(Q[s]) else 'normal'))

    plt.xticks([])
    plt.yticks([])

def choose_abs_greedy_action(state, Q, epsilon):
	action = None
	if np.random.uniform(0, 1) < epsilon:
		action = np.random.randint(env.action_space.n)
	else:
		action = np.argmax(Q[state,:])
	return action

def max_action_state(state, Q):
	action = np.argmax(Q[state,:])
	return Q[state, action]


def sarsa(env, alpha=0.1, gamma=0.9, epsilon=0.1, num_ep=int(1e4)):
    #Q = np.zeros((env.observation_space.n,  env.action_space.n))
    Q = np.random.rand(env.observation_space.n,  env.action_space.n)
    # TODO: implement the sarsa algorithm

    # This is some starting point performing random walks in the environment:
    for i in range(num_ep):
        s = env.reset()
        done = False
        a = choose_abs_greedy_action(s, Q, epsilon)
        while not done:
            s_, r, done, _ = env.step(a)
            a_ = choose_abs_greedy_action(s_, Q, epsilon)
            #update Q using sarsa
            Q[s, a] = Q[s, a] + alpha * (r + (gamma * Q[s_,a_]) - Q[s,a])
            s = s_
            a = a_
    return Q


def qlearning(env, alpha=0.1, gamma=0.9, epsilon=0.1, num_ep=int(1e4)):
    #Q = np.zeros((env.observation_space.n,  env.action_space.n))
    Q = np.random.rand(env.observation_space.n,  env.action_space.n)
    # TODO: implement the qlearning algorithm

    for i in range(num_ep):
        s = env.reset()
        done = False
        while not done:
            a = choose_abs_greedy_action(s, Q, epsilon)
            s_, r, done, _ = env.step(a)
            #update Q using Q learning

            Q[s, a] = Q[s, a] + alpha * (r+ ( gamma * max_action_state(s_, Q)) - Q[s,a] )
            s = s_
    return Q


env=gym.make('FrozenLake-v0')
#env=gym.make('FrozenLake-v0', is_slippery=False)
#env=gym.make('FrozenLake-v0', map_name="8x8")

print("Running sarsa...")
Q = sarsa(env)
plot_V(Q, env)
plot_Q(Q, env)
print_policy(Q, env)
plt.show()

print("Running qlearning")
Q = qlearning(env)
plot_V(Q, env)
plot_Q(Q, env)
print_policy(Q, env)
plt.show()

[(139, 'gym.make', 'gym.make', False, 'import gym\n'), (148, 'matplotlib.pyplot.show', 'plt.show', True, 'import matplotlib.pyplot as plt\n'), (155, 'matplotlib.pyplot.show', 'plt.show', True, 'import matplotlib.pyplot as plt\n'), (13, 'numpy.chararray', 'np.chararray', True, 'import numpy as np\n'), (26, 'matplotlib.pyplot.figure', 'plt.figure', True, 'import matplotlib.pyplot as plt\n'), (30, 'numpy.zeros', 'np.zeros', True, 'import numpy as np\n'), (36, 'matplotlib.pyplot.imshow', 'plt.imshow', True, 'import matplotlib.pyplot as plt\n'), (43, 'matplotlib.pyplot.xticks', 'plt.xticks', True, 'import matplotlib.pyplot as plt\n'), (44, 'matplotlib.pyplot.yticks', 'plt.yticks', True, 'import matplotlib.pyplot as plt\n'), (50, 'matplotlib.pyplot.figure', 'plt.figure', True, 'import matplotlib.pyplot as plt\n'), (57, 'numpy.array', 'np.array', True, 'import numpy as np\n'), (58, 'numpy.array', 'np.array', True, 'import numpy as np\n'), (59, 'numpy.array', 'np.array', True, 'import numpy as np\n'), (60, 'numpy.array', 'np.array', True, 'import numpy as np\n'), (65, 'matplotlib.colors.Normalize', 'colors.Normalize', False, 'from matplotlib import colors, patches\n'), (85, 'matplotlib.pyplot.xticks', 'plt.xticks', True, 'import matplotlib.pyplot as plt\n'), (86, 'matplotlib.pyplot.yticks', 'plt.yticks', True, 'import matplotlib.pyplot as plt\n'), (97, 'numpy.argmax', 'np.argmax', True, 'import numpy as np\n'), (103, 'numpy.random.rand', 'np.random.rand', True, 'import numpy as np\n'), (123, 'numpy.random.rand', 'np.random.rand', True, 'import numpy as np\n'), (16, 'numpy.unravel_index', 'np.unravel_index', True, 'import numpy as np\n'), (32, 'numpy.unravel_index', 'np.unravel_index', True, 'import numpy as np\n'), (33, 'numpy.max', 'np.max', True, 'import numpy as np\n'), (67, 'numpy.zeros', 'np.zeros', True, 'import numpy as np\n'), (71, 'numpy.unravel_index', 'np.unravel_index', True, 'import numpy as np\n'), (90, 'numpy.random.uniform', 'np.random.uniform', True, 'import numpy as np\n'), (91, 'numpy.random.randint', 'np.random.randint', True, 'import numpy as np\n'), (93, 'numpy.argmax', 'np.argmax', True, 'import numpy as np\n'), (17, 'numpy.argmax', 'np.argmax', True, 'import numpy as np\n'), (80, 'numpy.array', 'np.array', True, 'import numpy as np\n'), (83, 'numpy.max', 'np.max', True, 'import numpy as np\n')]