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, 4, 139, 29), 'gym.make', 'gym.make', ({(139, 13, 139, 28): '"""FrozenLake-v0"""'}, {}), "('FrozenLake-v0')", False, 'import gym\n'), ((148, 0, 148, 10), 'matplotlib.pyplot.show', 'plt.show', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((155, 0, 155, 10), 'matplotlib.pyplot.show', 'plt.show', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((13, 13, 13, 45), 'numpy.chararray', 'np.chararray', (), '', True, 'import numpy as np\n'), ((26, 10, 26, 22), 'matplotlib.pyplot.figure', 'plt.figure', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((30, 8, 30, 22), 'numpy.zeros', 'np.zeros', ({(30, 17, 30, 21): 'dims'}, {}), '(dims)', True, 'import numpy as np\n'), ((36, 4, 38, 56), 'matplotlib.pyplot.imshow', 'plt.imshow', (), '', True, 'import matplotlib.pyplot as plt\n'), ((43, 4, 43, 18), 'matplotlib.pyplot.xticks', 'plt.xticks', ({(43, 15, 43, 17): '[]'}, {}), '([])', True, 'import matplotlib.pyplot as plt\n'), ((44, 4, 44, 18), 'matplotlib.pyplot.yticks', 'plt.yticks', ({(44, 15, 44, 17): '[]'}, {}), '([])', True, 'import matplotlib.pyplot as plt\n'), ((50, 10, 50, 22), 'matplotlib.pyplot.figure', 'plt.figure', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((57, 9, 57, 46), 'numpy.array', 'np.array', ({(57, 18, 57, 45): '[[0, 1], [0.5, 0.5], [1, 1]]'}, {}), '([[0, 1], [0.5, 0.5], [1, 1]])', True, 'import numpy as np\n'), ((58, 11, 58, 48), 'numpy.array', 'np.array', ({(58, 20, 58, 47): '[[0, 0], [0.5, 0.5], [1, 0]]'}, {}), '([[0, 0], [0.5, 0.5], [1, 0]])', True, 'import numpy as np\n'), ((59, 11, 59, 48), 'numpy.array', 'np.array', ({(59, 20, 59, 47): '[[0, 0], [0.5, 0.5], [0, 1]]'}, {}), '([[0, 0], [0.5, 0.5], [0, 1]])', True, 'import numpy as np\n'), ((60, 12, 60, 49), 'numpy.array', 'np.array', ({(60, 21, 60, 48): '[[1, 0], [0.5, 0.5], [1, 1]]'}, {}), '([[1, 0], [0.5, 0.5], [1, 1]])', True, 'import numpy as np\n'), ((65, 11, 65, 44), 'matplotlib.colors.Normalize', 'colors.Normalize', (), '', False, 'from matplotlib import colors, patches\n'), ((85, 4, 85, 18), 'matplotlib.pyplot.xticks', 'plt.xticks', ({(85, 15, 85, 17): '[]'}, {}), '([])', True, 'import matplotlib.pyplot as plt\n'), ((86, 4, 86, 18), 'matplotlib.pyplot.yticks', 'plt.yticks', ({(86, 15, 86, 17): '[]'}, {}), '([])', True, 'import matplotlib.pyplot as plt\n'), ((97, 10, 97, 31), 'numpy.argmax', 'np.argmax', ({(97, 20, 97, 30): 'Q[(state), :]'}, {}), '(Q[(state), :])', True, 'import numpy as np\n'), ((103, 8, 103, 68), 'numpy.random.rand', 'np.random.rand', ({(103, 23, 103, 46): 'env.observation_space.n', (103, 49, 103, 67): 'env.action_space.n'}, {}), '(env.observation_space.n, env.action_space.n)', True, 'import numpy as np\n'), ((123, 8, 123, 68), 'numpy.random.rand', 'np.random.rand', ({(123, 23, 123, 46): 'env.observation_space.n', (123, 49, 123, 67): 'env.action_space.n'}, {}), '(env.observation_space.n, env.action_space.n)', True, 'import numpy as np\n'), ((16, 14, 16, 39), 'numpy.unravel_index', 'np.unravel_index', ({(16, 31, 16, 32): 's', (16, 34, 16, 38): 'dims'}, {}), '(s, dims)', True, 'import numpy as np\n'), ((32, 14, 32, 39), 'numpy.unravel_index', 'np.unravel_index', ({(32, 31, 32, 32): 's', (32, 34, 32, 38): 'dims'}, {}), '(s, dims)', True, 'import numpy as np\n'), ((33, 17, 33, 29), 'numpy.max', 'np.max', ({(33, 24, 33, 28): 'Q[s]'}, {}), '(Q[s])', True, 'import numpy as np\n'), ((67, 14, 67, 28), 'numpy.zeros', 'np.zeros', ({(67, 23, 67, 27): 'dims'}, {}), '(dims)', True, 'import numpy as np\n'), ((71, 14, 71, 39), 'numpy.unravel_index', 'np.unravel_index', ({(71, 31, 71, 32): 's', (71, 34, 71, 38): 'dims'}, {}), '(s, dims)', True, 'import numpy as np\n'), ((90, 4, 90, 27), 'numpy.random.uniform', 'np.random.uniform', ({(90, 22, 90, 23): '(0)', (90, 25, 90, 26): '(1)'}, {}), '(0, 1)', True, 'import numpy as np\n'), ((91, 11, 91, 48), 'numpy.random.randint', 'np.random.randint', ({(91, 29, 91, 47): 'env.action_space.n'}, {}), '(env.action_space.n)', True, 'import numpy as np\n'), ((93, 11, 93, 32), 'numpy.argmax', 'np.argmax', ({(93, 21, 93, 31): 'Q[(state), :]'}, {}), '(Q[(state), :])', True, 'import numpy as np\n'), ((17, 28, 17, 43), 'numpy.argmax', 'np.argmax', ({(17, 38, 17, 42): 'Q[s]'}, {}), '(Q[s])', True, 'import numpy as np\n'), ((80, 50, 80, 68), 'numpy.array', 'np.array', ({(80, 59, 80, 67): '[y, 3 - x]'}, {}), '([y, 3 - x])', True, 'import numpy as np\n'), ((83, 65, 83, 77), 'numpy.max', 'np.max', ({(83, 72, 83, 76): 'Q[s]'}, {}), '(Q[s])', True, 'import numpy as np\n')]