def animate_solution(town_map, path, save_anim=False, until_caught=False, gamma=0):
    col_map = {0: WHITE, 1: SEB_GREEN, 2: LIGHT_GREEN, (- 6): LIGHT_RED, (- 1): LIGHT_RED}
    (rows, cols) = town_map.shape
    fig = plt.figure(1, figsize=(cols, rows))
    ax = plt.gca()
    ax.set_title(('Policy simulation: $\\lambda$ = %0.1f' % gamma))
    ax.set_xticks([])
    ax.set_yticks([])
    colored_town_map = [[col_map[town_map[(j, i)]] for i in range(cols)] for j in range(rows)]
    fig = plt.figure(1, figsize=(cols, rows))
    grid = plt.table(cellText=None, cellColours=colored_town_map, cellLoc='center', loc=(0, 0), edges='closed')
    tc = grid.properties()['children']
    for cell in tc:
        cell.set_height((1.0 / rows))
        cell.set_width((1.0 / cols))
    path_robber = [tuple(p)[0:2] for p in path]
    path_police = [tuple(p)[2:] for p in path]
    for i in range(len(path_robber)):
        if (i == 0):
            grid.get_celld()[path_robber[i]].set_facecolor(LIGHT_ORANGE)
            grid.get_celld()[path_robber[i]].get_text().set_text('Robber')
            grid.get_celld()[path_police[i]].set_facecolor(LIGHT_RED)
            grid.get_celld()[path_police[i]].get_text().set_text('Police')
            if save_anim:
                plt.savefig(('optimal_policy_' + str(i)))
        else:
            if (until_caught and (path_robber[i] == path_police[i])):
                grid.get_celld()[path_robber[(i - 1)]].set_facecolor(col_map[town_map[path_robber[(i - 1)]]])
                grid.get_celld()[path_robber[(i - 1)]].get_text().set_text('')
                grid.get_celld()[path_police[(i - 1)]].set_facecolor(col_map[town_map[path_police[(i - 1)]]])
                grid.get_celld()[path_police[(i - 1)]].get_text().set_text('')
                grid.get_celld()[path_police[i]].set_facecolor(BUSTED_BLUE)
                grid.get_celld()[path_police[i]].get_text().set_text('BUSTED')
                print('BUSTED!!!', gamma)
                if save_anim:
                    plt.savefig((((str(gamma) + '_') + str(i)) + '.png'))
                break
            if save_anim:
                plt.savefig((((str(gamma) + '_') + str(i)) + '.png'))
            grid.get_celld()[path_robber[(i - 1)]].set_facecolor(col_map[town_map[path_robber[(i - 1)]]])
            grid.get_celld()[path_robber[(i - 1)]].get_text().set_text('')
            grid.get_celld()[path_police[(i - 1)]].set_facecolor(col_map[town_map[path_police[(i - 1)]]])
            grid.get_celld()[path_police[(i - 1)]].get_text().set_text('')
            grid.get_celld()[path_robber[i]].set_facecolor(LIGHT_ORANGE)
            grid.get_celld()[path_robber[i]].get_text().set_text('Robber')
            grid.get_celld()[path_police[i]].set_facecolor(LIGHT_RED)
            grid.get_celld()[path_police[i]].get_text().set_text('Police')
        grid.get_celld()[(0, 0)].get_text().set_text('SEB')
        grid.get_celld()[(0, 0)].get_text().set_color('white')
        grid.get_celld()[(0, 5)].get_text().set_text('SEB')
        grid.get_celld()[(0, 5)].get_text().set_color('white')
        grid.get_celld()[(2, 0)].get_text().set_text('SEB')
        grid.get_celld()[(2, 0)].get_text().set_color('white')
        grid.get_celld()[(2, 5)].get_text().set_text('SEB')
        grid.get_celld()[(2, 5)].get_text().set_color('white')
        plt.pause(0.7)
    plt.show()
def animate_solution(town_map, path, save_anim=False, until_caught=False, gamma=0):
    col_map = {0: WHITE, 1: SEB_GREEN, 2: LIGHT_GREEN, (- 6): LIGHT_RED, (- 1): LIGHT_RED}
    (rows, cols) = town_map.shape
    fig = plt.figure(1, figsize=(cols, rows))
    ax = plt.gca()
    ax.set_title(('Policy simulation: $\\lambda$ = %0.1f' % gamma))
    ax.set_xticks([])
    ax.set_yticks([])
    colored_town_map = [[col_map[town_map[(j, i)]] for i in range(cols)] for j in range(rows)]
    fig = plt.figure(1, figsize=(cols, rows))
    grid = plt.table(cellText=None, cellColours=colored_town_map, cellLoc='center', loc=(0, 0), edges='closed')
    tc = grid.properties()['children']
    for cell in tc:
        cell.set_height((1.0 / rows))
        cell.set_width((1.0 / cols))
    path_robber = [tuple(p)[0:2] for p in path]
    path_police = [tuple(p)[2:] for p in path]
    for i in range(len(path_robber)):
        if (i == 0):
            grid.get_celld()[path_robber[i]].set_facecolor(LIGHT_ORANGE)
            grid.get_celld()[path_robber[i]].get_text().set_text('Robber')
            grid.get_celld()[path_police[i]].set_facecolor(LIGHT_RED)
            grid.get_celld()[path_police[i]].get_text().set_text('Police')
            if save_anim:
                plt.savefig(('optimal_policy_' + str(i)))
        else:
            if (until_caught and (path_robber[i] == path_police[i])):
                grid.get_celld()[path_robber[(i - 1)]].set_facecolor(col_map[town_map[path_robber[(i - 1)]]])
                grid.get_celld()[path_robber[(i - 1)]].get_text().set_text()
                grid.get_celld()[path_police[(i - 1)]].set_facecolor(col_map[town_map[path_police[(i - 1)]]])
                grid.get_celld()[path_police[(i - 1)]].get_text().set_text()
                grid.get_celld()[path_police[i]].set_facecolor(BUSTED_BLUE)
                grid.get_celld()[path_police[i]].get_text().set_text('BUSTED')
                print('BUSTED!!!', gamma)
                if save_anim:
                    plt.savefig((((str(gamma) + '_') + str(i)) + '.png'))
                break
            if save_anim:
                plt.savefig((((str(gamma) + '_') + str(i)) + '.png'))
            grid.get_celld()[path_robber[(i - 1)]].set_facecolor(col_map[town_map[path_robber[(i - 1)]]])
            grid.get_celld()[path_robber[(i - 1)]].get_text().set_text()
            grid.get_celld()[path_police[(i - 1)]].set_facecolor(col_map[town_map[path_police[(i - 1)]]])
            grid.get_celld()[path_police[(i - 1)]].get_text().set_text()
            grid.get_celld()[path_robber[i]].set_facecolor(LIGHT_ORANGE)
            grid.get_celld()[path_robber[i]].get_text().set_text('Robber')
            grid.get_celld()[path_police[i]].set_facecolor(LIGHT_RED)
            grid.get_celld()[path_police[i]].get_text().set_text('Police')
        grid.get_celld()[(0, 0)].get_text().set_text('SEB')
        grid.get_celld()[(0, 0)].get_text().set_color('white')
        grid.get_celld()[(0, 5)].get_text().set_text('SEB')
        grid.get_celld()[(0, 5)].get_text().set_color('white')
        grid.get_celld()[(2, 0)].get_text().set_text('SEB')
        grid.get_celld()[(2, 0)].get_text().set_color('white')
        grid.get_celld()[(2, 5)].get_text().set_text('SEB')
        grid.get_celld()[(2, 5)].get_text().set_color('white')
        plt.pause(0.7)
    plt.show()