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

# Create model
def multilayer_perceptron(x, weights, biases):
    #somehow drl performs worse with complex layers
    # Hidden layer with RELU activation
    layer_1 = tf.matmul(x, weights['h1'])
    layer_1 = tf.nn.relu(layer_1)
    # Hidden layer with RELU activation
    layer_2 = tf.matmul(layer_1, weights['h2'])
    layer_2 = tf.nn.relu(layer_2)
    # Output layer with linear activation
    out_layer = tf.matmul(x, weights['out']) #+ biases['out']
    return out_layer




def build_model(env):
    tf.reset_default_graph()  # always call it first
    inputs1 = tf.placeholder(shape=[1, env.observation_space.n], dtype=tf.float32)  # tensor symbol for input state
    # Store layers weight & bias
    n_input = env.observation_space.n
    n_hidden_1 = 10
    n_hidden_2 = 10
    weights = {
        'h1': tf.Variable(tf.random_uniform([n_input, n_hidden_1], -0.01, 0.01)),
        'h2': tf.Variable(tf.random_uniform([n_hidden_1, n_hidden_2], -0.01, 0.01)),
        'out': tf.Variable(tf.random_uniform([n_input, env.action_space.n], -0.01, 0.01)),
    }
    biases = {
        'b1': tf.Variable(tf.random_uniform([n_hidden_1], -0.01, 0.01)),
        'b2': tf.Variable(tf.random_uniform([n_hidden_2], -0.01, 0.01)),
        'out': tf.Variable(tf.random_uniform([env.action_space.n], -0.01, 0.01))
    }
    Qout = multilayer_perceptron(inputs1, weights, biases)
    predict = tf.argmax(Qout, 1)  # given input_state, get best action id

    # Below we obtain the loss by taking the sum of squares difference between the target and prediction Q values.
    nextQ = tf.placeholder(shape=[1, env.action_space.n], dtype=tf.float32)  # tensor symbol for input Q
    loss = tf.reduce_sum(tf.square(nextQ - Qout))  # mean square loss function
    updateModel = tf.train.GradientDescentOptimizer(learning_rate=0.1).minimize(loss)
    return  inputs1, nextQ, Qout, predict, updateModel, weights, biases

def q_deep_learning(env, sess, inputs1, nextQ, Qout, predict, updateModel, weights, biases,
                    y = .99, e = 0.1, num_episodes = 2000):

    # create lists to contain total rewards and steps per episode
    jList = []
    rList = []
    numlost = numwin = 0
    for i in range(num_episodes):
        # Reset environment and get first new observation
        s = env.reset()
        rAll = 0
        # d = False
        j = 0
        stop = False

        # The Q-Network
        while j < 99:
            j += 1
            # Choose an action by greedily (with e chance of random action) from the Q-network
            # predict, Qout need input1 --> need feed_dict, s1 is id of state -->get onehot vector 1 at that id
            # get Qout as variable filled with values allQ --> Q values given current state
            a, allQ = sess.run([predict, Qout], feed_dict={inputs1: np.identity(env.observation_space.n)[s:s + 1]})
            if np.random.rand(1) < e:
                a[0] = env.action_space.sample()# random action still
            # Get new state and reward from environment
            s1, r, d, _ = env.step(a[0])# index 0 is index of max action
            # Obtain the Q' values by feeding the new state through our network
            # Q values given new state
            Q1 = sess.run(Qout, feed_dict={inputs1: np.identity(env.observation_space.n)[s1:s1 + 1]})
            # Obtain maxQ' and set our target value for chosen action.
            maxQ1 = np.max(Q1)
            targetQ = allQ
            targetQ[0, a[0]] = r + y * maxQ1 # assume model follow the rule --> next q values of (next state)
            # follow reward rules --> NN must predict q values of current state match this real value
            # if match --> Q converge --> reinforcement learning done!!!
            # Train our network using target and predicted Q values
            sess.run([updateModel], feed_dict={inputs1: np.identity(env.observation_space.n)[s:s + 1],
                                                          nextQ: targetQ})
            rAll += r
            s = s1
            if d:
                # Reduce chance of random action as we train the model.
                e = 1. / ((i / 50) + 10)
                if r == 0:
                    numlost += 1
                else:
                    numwin += 1
                stop = True
                break
        jList.append(j)
        rList.append(rAll/j)
        if not stop:
            numlost += 1

    print("Score over time: " + str(sum(rList) / num_episodes))
    print("Num win {} vs lost {} ".format(numwin / num_episodes, numlost / num_episodes))
    return updateModel


def q_table_learning(Q, env, lr = .5, y = .99, num_episodes = 2000):
    # Initialize table with all zeros

    # Set learning parameters

    # create lists to contain total rewards and steps per episode
    # jList = []
    rList = []
    numwin = 0
    numlost = 0
    for i in range(num_episodes):
        # Reset environment and get first new observation
        s = env.reset()
        rAll = 0
        j = 0
        stop=False
        # The Q-Table learning algorithm
        while j < 99:
            j += 1
            # Choose an action by greedily (with noise) picking from Q table
            a = np.argmax(Q[s, :] + np.random.randn(1, env.action_space.n) * (1. / (i + 1)))
            # Get new state and reward from environment
            s1, r, d, _ = env.step(a)
            # Update Q-Table with new knowledge
            Q[s, a] = Q[s, a] + lr * (r + y * np.max(Q[s1, :]) - Q[s, a])
            rAll += r
            s = s1
            if d:
                if r==0:
                    numlost+=1
                else:
                    numwin+=1
                stop=True
                break
        if not stop:
            numlost+=1
        rList.append(rAll/j)
        # jList.append(j)

    print("Score over time: " + str(sum(rList) / num_episodes))
    print("Num win {} vs lost {} ".format(numwin / num_episodes, numlost / num_episodes))
    # print("Step to win over time: " + str(sum(jList) / num_episodes))

def test_qtable():
    env = gym.make('FrozenLake-v0')
    numloop = 100
    Q = np.zeros([env.observation_space.n, env.action_space.n])
    print('start q table learning...')
    for i in range(numloop):
        print('============Loop: {} / {}================'.format(i, numloop))
        q_table_learning(Q, env)

def test_qdl():
    env = gym.make('FrozenLake-v0')
    inputs1, nextQ, Qout, predict, updateModel, weights, bias = build_model(env)
    init = tf.global_variables_initializer()  # auto init for all variable appear in tensorflow
    with tf.Session() as sess:# tensorflow run is based on session
        sess.run(init)
        numloop = 100
        print('start q deep learning...')
        for i in range(numloop):
            print('============Loop: {} / {}================'.format(i, numloop))
            q_deep_learning(env, sess, inputs1, nextQ, Qout, predict, updateModel, weights, bias)


if __name__ == '__main__':
    #test_qtable()
    test_qdl()


