import matplotlib.pyplot as plt
from matplotlib import colors
import numpy as np


def show_policy():                          # Visualising the hardcoded policy

    with open("Option1.txt") as textFile:
        grid = [line.split() for line in textFile]
    grid = np.array(grid).astype(np.float)

    # with open("Option2.txt") as textFile:
    #     grid = [line.split() for line in textFile]
    # grid = np.array(grid).astype(np.float)

    cmap = colors.ListedColormap([(0, 0, 0), (1, 1, 1), (1, 1, 1), (1, 1, 1), (1, 1, 1)])

    plt.figure(figsize=(13, 13))
    plt.pcolor(grid, cmap=cmap, edgecolors='k', linewidths=5)

    ax = plt.axes()
    for i in range(0, 13):
        for j in range(0, 13):

            if grid[j][i] == -1:
                continue

            if grid[j][i] == 2:
                ax.arrow(i + 0.5, j + 0.2, 0, 0.4, head_width=0.25, head_length=0.25, fc='g', ec='g')
            elif grid[j][i] == 1:
                ax.arrow(i + 0.2, j + 0.5, 0.4, 0, head_width=0.25, head_length=0.25, fc='g', ec='g')
            elif grid[j][i] == 0:
                ax.arrow(i + 0.5, j + 0.8, 0, -0.4, head_width=0.25, head_length=0.25, fc='g', ec='g')
            else:
                ax.arrow(i + 0.8, j + 0.5, -0.4, 0, head_width=0.25, head_length=0.25, fc='g', ec='g')

    plt.gca().invert_yaxis()
    plt.show()

show_policy()

