
"""
Many thanks for the contribution of Dr. Mofan Zhou's Reinforcement Learning tutorial
The A2C part of this project base on the github Repositories of Dr. Zhou.
The original information of Dr. Zhou's tutorials as follows.
"""

"""
Actor-Critic with continuous action using TD-error as the Advantage, Reinforcement Learning.

The Pendulum example (based on https://github.com/dennybritz/reinforcement-learning/blob/master/PolicyGradient/Continuous%20MountainCar%20Actor%20Critic%20Solution.ipynb)

Cannot converge!!! oscillate!!!

View more on my tutorial page: https://morvanzhou.github.io/tutorials/

Using:
tensorflow r1.3
gym 0.8.0
"""

import tensorflow as tf
import numpy as np
import gym
import ARIS_ENV as arenv
import globe
import matplotlib.pyplot as plt
import sys

np.random.seed(2)
tf.set_random_seed(2)  # reproducible

##########################################################################
action_epsilon = 1
action_epsilon_decay = 1/2000
MAX_EP_STEPS = 600

idx = sys.argv[1]
globe._init()
arenv._init(globe, True, MAX_EP_STEPS, "Dataset/Distance_"+str(idx)+".csv")


class Actor(object):
    def __init__(self, sess, n_features, action_bound, lr=0.0001):
        self.sess = sess

        self.s = tf.placeholder(tf.float32, [1, n_features], "state")
        self.a = tf.placeholder(tf.float32, None, name="act")
        self.td_error = tf.placeholder(tf.float32, None, name="td_error")  # TD_error

        l1 = tf.layers.dense(
            inputs=self.s,
            units=30,  # number of hidden units
            activation=tf.nn.relu,
            kernel_initializer=tf.random_normal_initializer(0., .1),  # weights
            bias_initializer=tf.constant_initializer(0.1),  # biases
            name='l1'
        )

        mu = tf.layers.dense(
            inputs=l1,
            units=1,  # number of hidden units
            activation=tf.nn.tanh,
            kernel_initializer=tf.random_normal_initializer(0., .1),  # weights
            bias_initializer=tf.constant_initializer(0.1),  # biases
            name='mu'
        )

        sigma = tf.layers.dense(
            inputs=l1,
            units=1,  # output units
            activation=tf.nn.softplus,  # get action probabilities
            kernel_initializer=tf.random_normal_initializer(0., .1),  # weights
            bias_initializer=tf.constant_initializer(1.),  # biases
            name='sigma'
        )
        global_step = tf.Variable(0, trainable=False)
        # self.e = epsilon = tf.train.exponential_decay(2., global_step, 1000, 0.9)
        self.mu, self.sigma = tf.squeeze(mu*2), tf.squeeze(sigma+0.1)
        self.normal_dist = tf.distributions.Normal(self.mu, self.sigma)

        self.action = tf.clip_by_value(self.normal_dist.sample(2), action_bound[0], action_bound[1])

        with tf.name_scope('exp_v'):
            log_prob = self.normal_dist.log_prob(self.a)  # loss without advantage
            self.exp_v = log_prob * self.td_error  # advantage (TD_error) guided loss
            # Add cross entropy cost to encourage exploration
            self.exp_v += 0.01*self.normal_dist.entropy()

        with tf.name_scope('train'):
            self.train_op = tf.train.AdamOptimizer(lr).minimize(-self.exp_v, global_step)    # min(v) = max(-v)

    def learn(self, s, a, td):
        s = s[np.newaxis, :]
        feed_dict = {self.s: s, self.a: a, self.td_error: td}
        _, exp_v = self.sess.run([self.train_op, self.exp_v], feed_dict)
        return exp_v

    def choose_action(self, s):
        # s = s[np.newaxis, :]
        # return self.sess.run(self.action, {self.s: s})  # get probabilities for all actions
        return self.sess.run(self.action, {self.s: s[np.newaxis, :]})[0]


class Critic(object):
    def __init__(self, sess, n_features, lr=0.01):
        self.sess = sess
        with tf.name_scope('inputs'):
            self.s = tf.placeholder(tf.float32, [1, n_features], "state")
            self.v_ = tf.placeholder(tf.float32, [1, 1], name="v_next")
            self.r = tf.placeholder(tf.float32, name='r')

        with tf.variable_scope('Critic'):
            l1 = tf.layers.dense(
                inputs=self.s,
                units=30,  # number of hidden units
                activation=tf.nn.relu,
                kernel_initializer=tf.random_normal_initializer(0., .1),  # weights
                bias_initializer=tf.constant_initializer(0.1),  # biases
                name='l1'
            )

            self.v = tf.layers.dense(
                inputs=l1,
                units=1,  # output units
                activation=None,
                kernel_initializer=tf.random_normal_initializer(0., .1),  # weights
                bias_initializer=tf.constant_initializer(0.1),  # biases
                name='V'
            )

        with tf.variable_scope('squared_TD_error'):
            self.td_error = tf.reduce_mean(self.r + GAMMA * self.v_ - self.v)
            self.loss = tf.square(self.td_error)    # TD_error = (r+gamma*V_next) - V_eval
        with tf.variable_scope('train'):
            self.train_op = tf.train.AdamOptimizer(lr).minimize(self.loss)

    def learn(self, s, r, s_):
        s, s_ = s[np.newaxis, :], s_[np.newaxis, :]

        v_ = self.sess.run(self.v, {self.s: s_})
        td_error, _ = self.sess.run([self.td_error, self.train_op],
                                          {self.s: s, self.v_: v_, self.r: r})
        return td_error

def ou_noise(x, mu, theta, sigma):
    return theta * (mu - x) + sigma * np.random.randn(1)

def plot(idx, frame_idx, rewards):
    plt.figure()
    plt.title('Dataset %i, frame %s. reward: %s' % (int(idx), frame_idx, rewards[-1]))
    # print('frame %s. reward: %s' % (frame_idx, rewards[-1]))
    plt.plot(rewards)
    plt.show()
    plt.savefig('A2C_Result/Dataset_'+str(idx)+'.png', format='png')
    plt.close()


OUTPUT_GRAPH = False
MAX_EPISODE = 500
MAX_EP_STEPS = 600
DISPLAY_REWARD_THRESHOLD = 5  # renders environment if total episode reward is greater then this threshold
RENDER = False  # rendering wastes time
GAMMA = 0.99
LR_A = 0.001    # learning rate for actor
LR_C = 0.01     # learning rate for critic


observation_space = arenv._observation_space()
action_space = arenv._action_space()
N_S = observation_space.shape[0]
A_BOUND = [0.5, 0.5]#env.action_space.high

sess = tf.Session()

actor = Actor(sess, n_features=N_S, lr=LR_A, action_bound=A_BOUND)
critic = Critic(sess, n_features=N_S, lr=LR_C)

sess.run(tf.global_variables_initializer())

rewards = []
var = 1  # control exploration

if OUTPUT_GRAPH:
    tf.summary.FileWriter("logs/", sess.graph)

for i_episode in range(MAX_EPISODE):
    s = arenv.reset()
    t = 0
    ep_rs = []
    while True:

        a = actor.choose_action(s)
        # noise
        action_epsilon -= action_epsilon_decay
        noise = max(action_epsilon, 0) * ou_noise(a, 0.0, 0.6, 0.30)
        a = np.clip(np.random.normal(a+noise, var, size=2), 0.01, 0.99)    # add randomness to action selection for exploration
        # print(a)
        var *= .9995

        s_, r, done = arenv.Step(a)

        td_error = critic.learn(s, r/10, s_)  # gradient = grad[r + gamma * V(s_) - V(s)]
        actor.learn(s, a, td_error)  # true_gradient = grad[logPi(s,a) * td_error]

        s = s_
        t += 1
        ep_rs.append(r)
        if t > MAX_EP_STEPS-1:
            ep_rs_sum = sum(ep_rs)

            rewards.append(ep_rs_sum)
            np.savetxt("A2C_Result/Dataset_"+str(idx)+"_Rewards.csv", rewards, delimiter=',')

            print("episode:", i_episode, "  reward:", rewards[-1])
            break

    if i_episode > 1:
        plot(idx, i_episode, rewards)
