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()

[((4710, 4735), 'gym.make', 'gym.make', (['"""FrozenLake-v0"""'], {}), "('FrozenLake-v0')\n", (4718, 4735), False, 'import gym\n'), ((4926, 4936), 'matplotlib.pyplot.show', 'plt.show', ([], {}), '()\n', (4934, 4936), True, 'import matplotlib.pyplot as plt\n'), ((5035, 5045), 'matplotlib.pyplot.show', 'plt.show', ([], {}), '()\n', (5043, 5045), True, 'import matplotlib.pyplot as plt\n'), ((330, 362), 'numpy.chararray', 'np.chararray', (['dims'], {'unicode': '(True)'}), '(dims, unicode=True)\n', (342, 362), True, 'import numpy as np\n'), ((780, 792), 'matplotlib.pyplot.figure', 'plt.figure', ([], {}), '()\n', (790, 792), True, 'import matplotlib.pyplot as plt\n'), ((882, 896), 'numpy.zeros', 'np.zeros', (['dims'], {}), '(dims)\n', (890, 896), True, 'import numpy as np\n'), ((1063, 1191), 'matplotlib.pyplot.imshow', 'plt.imshow', (['V'], {'origin': '"""upper"""', 'extent': '[0, dims[0], 0, dims[1]]', 'vmin': '(0.0)', 'vmax': '(0.6)', 'cmap': 'plt.cm.RdYlGn', 'interpolation': '"""none"""'}), "(V, origin='upper', extent=[0, dims[0], 0, dims[1]], vmin=0.0,\n    vmax=0.6, cmap=plt.cm.RdYlGn, interpolation='none')\n", (1073, 1191), True, 'import matplotlib.pyplot as plt\n'), ((1431, 1445), 'matplotlib.pyplot.xticks', 'plt.xticks', (['[]'], {}), '([])\n', (1441, 1445), True, 'import matplotlib.pyplot as plt\n'), ((1450, 1464), 'matplotlib.pyplot.yticks', 'plt.yticks', (['[]'], {}), '([])\n', (1460, 1464), True, 'import matplotlib.pyplot as plt\n'), ((1601, 1613), 'matplotlib.pyplot.figure', 'plt.figure', ([], {}), '()\n', (1611, 1613), True, 'import matplotlib.pyplot as plt\n'), ((1725, 1763), 'numpy.array', 'np.array', (['[[0, 1], [0.5, 0.5], [1, 1]]'], {}), '([[0, 1], [0.5, 0.5], [1, 1]])\n', (1733, 1763), True, 'import numpy as np\n'), ((1774, 1812), 'numpy.array', 'np.array', (['[[0, 0], [0.5, 0.5], [1, 0]]'], {}), '([[0, 0], [0.5, 0.5], [1, 0]])\n', (1782, 1812), True, 'import numpy as np\n'), ((1823, 1861), 'numpy.array', 'np.array', (['[[0, 0], [0.5, 0.5], [0, 1]]'], {}), '([[0, 0], [0.5, 0.5], [0, 1]])\n', (1831, 1861), True, 'import numpy as np\n'), ((1873, 1911), 'numpy.array', 'np.array', (['[[1, 0], [0.5, 0.5], [1, 1]]'], {}), '([[1, 0], [0.5, 0.5], [1, 1]])\n', (1881, 1911), True, 'import numpy as np\n'), ((2045, 2081), 'matplotlib.colors.Normalize', 'colors.Normalize', ([], {'vmin': '(0.0)', 'vmax': '(0.6)'}), '(vmin=0.0, vmax=0.6)\n', (2061, 2081), False, 'from matplotlib import colors, patches\n'), ((3027, 3041), 'matplotlib.pyplot.xticks', 'plt.xticks', (['[]'], {}), '([])\n', (3037, 3041), True, 'import matplotlib.pyplot as plt\n'), ((3046, 3060), 'matplotlib.pyplot.yticks', 'plt.yticks', (['[]'], {}), '([])\n', (3056, 3060), True, 'import matplotlib.pyplot as plt\n'), ((3312, 3336), 'numpy.argmax', 'np.argmax', (['Q[(state), :]'], {}), '(Q[(state), :])\n', (3321, 3336), True, 'import numpy as np\n'), ((3503, 3562), 'numpy.random.rand', 'np.random.rand', (['env.observation_space.n', 'env.action_space.n'], {}), '(env.observation_space.n, env.action_space.n)\n', (3517, 3562), True, 'import numpy as np\n'), ((4240, 4299), 'numpy.random.rand', 'np.random.rand', (['env.observation_space.n', 'env.action_space.n'], {}), '(env.observation_space.n, env.action_space.n)\n', (4254, 4299), True, 'import numpy as np\n'), ((425, 450), 'numpy.unravel_index', 'np.unravel_index', (['s', 'dims'], {}), '(s, dims)\n', (441, 450), True, 'import numpy as np\n'), ((939, 964), 'numpy.unravel_index', 'np.unravel_index', (['s', 'dims'], {}), '(s, dims)\n', (955, 964), True, 'import numpy as np\n'), ((982, 994), 'numpy.max', 'np.max', (['Q[s]'], {}), '(Q[s])\n', (988, 994), True, 'import numpy as np\n'), ((2098, 2112), 'numpy.zeros', 'np.zeros', (['dims'], {}), '(dims)\n', (2106, 2112), True, 'import numpy as np\n'), ((2302, 2327), 'numpy.unravel_index', 'np.unravel_index', (['s', 'dims'], {}), '(s, dims)\n', (2318, 2327), True, 'import numpy as np\n'), ((3130, 3153), 'numpy.random.uniform', 'np.random.uniform', (['(0)', '(1)'], {}), '(0, 1)\n', (3147, 3153), True, 'import numpy as np\n'), ((3176, 3213), 'numpy.random.randint', 'np.random.randint', (['env.action_space.n'], {}), '(env.action_space.n)\n', (3193, 3213), True, 'import numpy as np\n'), ((3232, 3256), 'numpy.argmax', 'np.argmax', (['Q[(state), :]'], {}), '(Q[(state), :])\n', (3241, 3256), True, 'import numpy as np\n'), ((479, 494), 'numpy.argmax', 'np.argmax', (['Q[s]'], {}), '(Q[s])\n', (488, 494), True, 'import numpy as np\n'), ((2723, 2743), 'numpy.array', 'np.array', (['[y, 3 - x]'], {}), '([y, 3 - x])\n', (2731, 2743), True, 'import numpy as np\n'), ((2993, 3005), 'numpy.max', 'np.max', (['Q[s]'], {}), '(Q[s])\n', (2999, 3005), True, 'import numpy as np\n')]