#!/usr/bin/env python

# -*- encoding: utf-8

'''
_____.___._______________  __.____ __________    _________   ___ ___    _____  .___ 
\__  |   |\_   _____/    |/ _|    |   \      \   \_   ___ \ /   |   \  /  _  \ |   |
 /   |   | |    __)_|      < |    |   /   |   \  /    \  \//    ~    \/  /_\  \|   |
 \____   | |        \    |  \|    |  /    |    \ \     \___\    Y    /    |    \   |
 / ______|/_______  /____|__ \______/\____|__  /  \______  /\___|_  /\____|__  /___|
 \/               \/        \/               \/          \/       \/         \/     
 

@author: Yekun Chai
@license: IACAS
@contact: chaiyekun@gmail.com
@file: AC_CartPole.py
@time: 3/27/19 5:06 PM
@desc：       
               
'''

import numpy  as np
import tensorflow as tf
import gym

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

OUTPUT_GRAPH = False
MAX_EPISODE = 3000
DISPLAY_REWARD_THRESHOLD = 200  # render the env if total episode reward is greater than this threshold
MAX_EP_STEPS = 1000  # maximum time step in one episode
RENDER = False
GAMMA = .9  # reward discount in TD error
LR_A = .001  # lr for actor
LR_C = .01  # lr for critic

env = gym.make("CartPole-v0")
env.seed(1)
env = env.unwrapped

N_F = env.observation_space.shape[0]
N_A = env.action_space.n


class Actor:
    def __init__(self, sess, n_features, n_actions, lr=.001):
        self.sess = sess
        self.s = tf.placeholder(tf.float32, [1, n_features], "state")
        self.a = tf.placeholder(tf.int32, None, "act")
        self.td_error = tf.placeholder(tf.float32, None, "td_error")

        with tf.variable_scope("Actor"):
            l1 = tf.layers.dense(inputs=self.s, units=20, activation=tf.nn.relu,
                                 kernel_initializer=tf.random_normal_initializer(0, .1),
                                 bias_initializer=tf.constant_initializer(.1), name="l1")

            self.acts_prob = tf.layers.dense(inputs=l1, units=n_actions, activation=tf.nn.softmax,
                                             kernel_initializer=tf.random_normal_initializer(0., .1),
                                             bias_initializer=tf.constant_initializer(.1), name="acts_prob")

        with tf.variable_scope("exp_v"):
            log_prob = tf.log(self.acts_prob[0, self.a])
            self.exp_v = tf.reduce_mean(log_prob * self.td_error)

        with tf.variable_scope("train"):
            self.train_op = tf.train.AdamOptimizer(lr).minimize(-self.exp_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, :]
        probs = self.sess.run(self.acts_prob, {self.s: s})
        return np.random.choice(np.arange(probs.shape[1]), p=probs.ravel())


class Critic:
    def __init__(self, sess, n_features, lr=0.01):
        self.sess = sess

        self.s = tf.placeholder(tf.float32, [1, n_features], "state")
        self.v_ = tf.placeholder(tf.float32, [1, 1], "v_next")
        self.r = tf.placeholder(tf.float32, None, "r")

        with tf.variable_scope("Critic"):
            l1 = tf.layers.dense(inputs=self.s, units=20, activation=tf.nn.relu,
                                 kernel_initializer=tf.random_normal_initializer(0, .1),
                                 bias_initializer=tf.constant_initializer(.1), name="l1")

            self.v = tf.layers.dense(inputs=l1, units=1, activation=None,
                                     kernel_initializer=tf.random_normal_initializer(0, .1),
                                     bias_initializer=tf.constant_initializer(.1), name="V")

        with tf.variable_scope("squared_TD_error"):
            self.td_error = self.r + GAMMA * self.v_ - self.v
            self.loss = tf.square(self.td_error)
        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


if __name__ == '__main__':
    sess = tf.Session()
    actor = Actor(sess, n_features=N_F, n_actions=N_A, lr=LR_A)
    critic = Critic(sess, n_features=N_F, lr=LR_C) # the teacher learn faster than the actor

    sess.run(tf.global_variables_initializer())

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

    for episode in range(MAX_EPISODE):
        s = env.reset()
        t = 0
        track_r = []
        while True:
            if RENDER: env.render()
            a = actor.choose_action(s)

            s_, r, done, info = env.step(a)

            if done: r = -20

            track_r.append(r)

            td_error = critic.learn(s, r, 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
            if done or t >= MAX_EP_STEPS:
                ep_rs_sum = sum(track_r)

                if "running_reward" not in globals():
                    running_reward = ep_rs_sum
                else:
                    running_reward = running_reward * .95 + ep_rs_sum * .05
                if running_reward > DISPLAY_REWARD_THRESHOLD:
                    RENDER = True
                print("episode:", episode, " reward:", int(running_reward))
                break
