import tensorflow as tf # neural network for function approximation
import gym # environment
import numpy as np # matrix operation and math functions
from gym import wrappers
import gym_morph # customized environment for cart-pole
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import time

start_time = time.time()
max_test = 10

for test_num in range(1, max_test):

    # Hyperparameters
    RANDOM_NUMBER_SEED = test_num

    ENVIRONMENT1 = "morph-v0"

    MAX_EPISODES = 5000 # number of episodes
    EPISODE_LENGTH = 200 # single episode length

    HIDDEN_SIZE = 16
    DISPLAY_WEIGHTS = False # Help debug weight update

    gamma = 0.99  # Discount per step

    RENDER = False  # Render the cart-pole system
    VIDEO_INTERVAL = 100  # Generate a video at this interval

    CONSECUTIVE_TARGET = 100 # Including previous 100 rewards

    CONST_LR = False # Constant or decaying learing rate
    # Constant learning rate
    const_learning_rate_in = 0.006
    # Decay learning rate
    start_learning_rate_in = 0.005
    decay_steps_in = 100
    decay_rate_in = 0.92

    DIR_PATH_SAVEFIG = "/root/cartpole_plot/"

    if CONST_LR:
        learning_rate = const_learning_rate_in

        file_name_savefig = "el" + str(EPISODE_LENGTH) \
        + "_hn" + str(HIDDEN_SIZE) \
        + "_clr" + str(learning_rate).replace(".", "p") \
        + "_test" + str(test_num) \
        + ".png"
    else:
        start_learning_rate = start_learning_rate_in
        decay_steps = decay_steps_in
        decay_rate = decay_rate_in

        file_name_savefig = "el" + str(EPISODE_LENGTH) \
        + "_hn" + str(HIDDEN_SIZE) \
        + "_dlr_slr" + str(start_learning_rate).replace(".", "p") \
        + "_ds" + str(decay_steps) \
        + "_dr" + str(decay_rate).replace(".", "p") \
        + "_test" + str(test_num) \
        + ".png"

    env = gym.make(ENVIRONMENT1)

    env.seed(RANDOM_NUMBER_SEED)
    np.random.seed(RANDOM_NUMBER_SEED)
    tf.set_random_seed(RANDOM_NUMBER_SEED)

    # Input and output sizes
    input_size = 4
    output_size = 2
    # input_size = env.observation_space.shape[0]
    # try:
    #     output_size = env.action_space.shape[0]
    # except AttributeError:
    #     output_size = env.action_space.n

    # Tensorflow network setup
    x = tf.placeholder(tf.float32, shape=(None, input_size))
    y = tf.placeholder(tf.float32, shape=(None, 1))

    if not CONST_LR:
        # decay learning rate
        global_step = tf.Variable(0, trainable=False)
        learning_rate = tf.train.exponential_decay(start_learning_rate, global_step, decay_steps, decay_rate, staircase=False)

    expected_returns = tf.placeholder(tf.float32, shape=(None, 1))

    # Xavier (2010) weights initializer for uniform distribution:
    # x = sqrt(6. / (in + out)); [-x, x]
    w_init = tf.contrib.layers.xavier_initializer()

    hidden_W = tf.get_variable("W1", shape=[input_size, HIDDEN_SIZE],
    initializer=w_init)
    hidden_B = tf.Variable(tf.zeros(HIDDEN_SIZE))
    dist_W = tf.get_variable("W2", shape=[HIDDEN_SIZE, output_size],
    initializer=w_init)
    dist_B = tf.Variable(tf.zeros(output_size))
    hidden = tf.nn.elu(tf.matmul(x, hidden_W) + hidden_B)
    dist = tf.tanh(tf.matmul(hidden, dist_W) + dist_B)

    dist_soft = tf.nn.log_softmax(dist)
    dist_in = tf.matmul(dist_soft, tf.Variable([[1.], [0.]]))
    pi = tf.contrib.distributions.Bernoulli(dist_in)
    pi_sample = pi.sample()
    log_pi = pi.log_prob(y)

    if CONST_LR:
        optimizer = tf.train.RMSPropOptimizer(learning_rate)
        train = optimizer.minimize(-1.0 * expected_returns * log_pi)
    else:
        optimizer = tf.train.RMSPropOptimizer(learning_rate)
        train = optimizer.minimize(-1.0 * expected_returns * log_pi, global_step=global_step)

    # saver = tf.train.Saver()

    # Create and initialize a session
    sess = tf.Session()
    sess.run(tf.global_variables_initializer())

    def run_episode(environment, ep, render=False):
        raw_reward = 0
        discounted_reward = 0
        cumulative_reward = []
        discount = 1.0
        states = []
        actions = []
        obs = environment.reset()
        done = False
        while not done:
            states.append(obs)
            cumulative_reward.append(discounted_reward)
            if render and ((ep % VIDEO_INTERVAL) == 0):
                environment.render()
            action = sess.run(pi_sample, feed_dict={x: [obs]})[0]
            actions.append(action)
            obs, reward, done, info = env.step(action[0])
            raw_reward += reward
            if reward > 0:
                discounted_reward += reward * discount
            else:
                discounted_reward += reward
            discount *= gamma
        return raw_reward, discounted_reward, cumulative_reward, states, actions


    def display_weights(session):
        w1 = session.run(hidden_W)
        b1 = session.run(hidden_B)
        w2 = session.run(dist_W)
        b2 = session.run(dist_B)
        print(w1, b1, w2, b2)

    returns = []
    mean_returns = []
    for ep in range(MAX_EPISODES):
        raw_G, discounted_G, cumulative_G, ep_states, ep_actions = \
            run_episode(env, ep, RENDER)
        expected_R = np.transpose([discounted_G - np.array(cumulative_G)])
        sess.run(train, feed_dict={x: ep_states, y: ep_actions,
                                   expected_returns: expected_R})
        if DISPLAY_WEIGHTS:
            display_weights(sess)

        returns.append(raw_G)
        running_returns = returns[max(0, ep-CONSECUTIVE_TARGET):(ep+1)]
        mean_return = np.mean(running_returns)
        mean_returns.append(mean_return)
        if CONST_LR:
            msg = "Test: {}, Episode: {}, Time: {}, Learning rate: {}, Return: {}, Last {} returns mean: {}"
            msg = msg.format(test_num, ep+1, time.strftime('%H:%M:%S', time.gmtime(time.time()-start_time)), learning_rate, raw_G, CONSECUTIVE_TARGET, mean_return)
            print(msg)
        else:
            msg = "Test: {}, Episode: {}, Time: {}, Learning rate: {}, Return: {}, Last {} returns mean: {}"
            msg = msg.format(test_num, ep+1, time.strftime('%H:%M:%S', time.gmtime(time.time()-start_time)), sess.run(learning_rate), raw_G, CONSECUTIVE_TARGET, mean_return)
            print(msg)

    env.close() # close openai gym environment
    tf.reset_default_graph() # clear tensorflow graph

    # Plot
    # plt.style.use('ggplot')
    plt.style.use('dark_background')
    episodes_plot = np.arange(MAX_EPISODES)
    fig = plt.figure()
    ax = fig.add_subplot(111)
    fig.subplots_adjust(top=0.85)
    if CONST_LR:
        ax.set_title("The Cart-Pole Problem Test %i \n \
        Episode Length: %i   \
        Discount Factor: %.2f \n \
        Number of Hidden Neuron: %i   \
        Constant Learning Rate: %.5f" % (test_num, EPISODE_LENGTH, gamma, HIDDEN_SIZE, learning_rate))
    else:
        ax.set_title("The Cart-Pole Problem Test %i \n \
        EpisodeLength: %i  DiscountFactor: %.2f  NumHiddenNeuron: %i \n \
        Decay Learning Rate: (start: %.5f, steps: %i, rate: %.2f)" % (test_num, EPISODE_LENGTH, gamma, HIDDEN_SIZE, start_learning_rate, decay_steps, decay_rate))
    ax.set_xlabel("Episode")
    ax.set_ylabel("Return")
    ax.set_ylim((0, EPISODE_LENGTH))
    ax.grid(linestyle='--')
    ax.plot(episodes_plot, returns, label='Instant return')
    ax.plot(episodes_plot, mean_returns, label='Averaged return')
    legend = ax.legend(loc='best', shadow=True)
    fig.savefig(DIR_PATH_SAVEFIG + file_name_savefig, dpi=500)
    # plt.show()
