import gym
import numpy as np
import random
import tensorflow as tf
import matplotlib.pyplot as plt

#Load the environment
env = gym.make('FrozenLake-v0')

#Implement our neural net
tf.reset_default_graph()

#Create the feed-forward portion of the nural net. These are used to choose actions
inputs1 = tf.placeholder(shape=[1,16], dtype=tf.float32) #Note: 16 inputs for the 16 possible actions (4 actions from each of 4 states)
W = tf.Variable(tf.random_uniform([16,4], 0, 0.01)) #Create a weight matrix with some initial noise
Qout = tf.matmul(inputs1, W) #We output a vector with 4 values: 1 for each action
predict = tf.argmax(Qout,1) #Choose the highest weighted action

#Below we obtaint he loss by taking the sum of squares difference between the target and the prediction Q values

nextQ = tf.placeholder(shape=[1,4], dtype=tf.float32)
loss = tf.reduce_sum(tf.square(nextQ - Qout))
trainer = tf.train.GradientDescentOptimizer(learning_rate=0.1) #The learning rate is a hyper parameter we need to tune
updateModel = trainer.minimize(loss)

#Now that we have the network, let's train it!

init = tf.initialize_all_variables()

#Set our learning parameters
y = .99
e = 0.1
num_episodes = 2000
#create lists to contain total rewards and steps per episode
jList = []
rList = []
with tf.Session() as sess:
    sess.run(init)
    for i in range(num_episodes):
        #Reset the env and get first new observation
        s = env.reset()
        rAll = 0
        d = False
        j =0
        #The Q-network
        while j < 99:
            j += 1
            #Choose an action by greedily (with "e" chance of random action) from the Q-network
            #Note that this acts like Simulated Annealing
            a, allQ = sess.run([predict, Qout], feed_dict={inputs1:np.identity(16)[s:s+1]})
            if np.random.rand(1) < e:
                a[0] = env.action_space.sample()
            #Get new state and reward from env
            s1, r, d, _ = env.step(a[0])
            #Obtain the Q values by feeding the new state through our network
            Q1 = sess.run(Qout, feed_dict={inputs1:np.identity(16)[s1:s1+1]})
            #Obtain maxQ and set our target value for chose action
            maxQ1 = np.max(Q1)
            targetQ = allQ
            targetQ[0,a[0]] = r + y*maxQ1
            #train our network using the target and predicted Q values
            _,W1 = sess.run([updateModel, W], feed_dict={inputs1:np.identity(16)[s:s+1], nextQ:targetQ})
            rAll += r
            s = s1
            if d == True:
                #Reduce the chance of random action as we train the model
                e = 1./((i/50) + 10)
                break
        jList.append(j)
        rList.append(rAll)
print("Percent of successful episodes: " + str(sum(rList)/num_episodes) + "%")
plt.plot(rList)
plt.show()
plt.gcf().clear()
plt.plot(jList)
plt.show()

