"""
State-Value Function

Written by Patrick Coady (pat-coady.github.io)

Modified by Tin-Yin Lai (wu6u3) into asynchronous version
"""

import tensorflow as tf
import numpy as np
from sklearn.utils import shuffle
#import os

class NNValueFunction(object):
    """ NN-based state-value function """
    def __init__(self, obs_dim, hid1_mult, thread_idx, shared_nn):
        """
        Args:
            obs_dim: number of dimensions in observation vector (int)
            hid1_mult: size of first hidden layer, multiplier of obs_dim
        """
        self.replay_buffer_x = None
        self.replay_buffer_y = None
        self.obs_dim = obs_dim
        self.hid1_mult = hid1_mult
        self.epochs = 10
        self.lr = None  # learning rate set in _build_graph()
        self._thread_idx=thread_idx # -1 for global
        self._scope_name = "nn_net_"+str(self._thread_idx)
        self._build_graph()
        #self.sess = tf.Session(graph=self.g)
        #self.sess.run(self.init)

        var_refs = [v._ref() for v in self.get_vars()]
        self.gradients = tf.gradients(
            self.loss, var_refs,
            gate_gradients=False,
            aggregation_method=None,
            colocate_gradients_with_ops=False)
        self.apply_gradients=None
        self.sync = self.sync_from(shared_nn)

        #self. global_fit = self.fit_for_global(x=None, y=None, logger=None)

    def _build_graph(self):
        """ Construct TensorFlow graph, including loss function, init op and train op """
        with tf.variable_scope(self._scope_name) as scope:
            self.obs_ph = tf.placeholder(tf.float32, (None, self.obs_dim), 'obs_valfunc')
            self.val_ph = tf.placeholder(tf.float32, (None,), 'val_valfunc')
            # hid1 layer size is 10x obs_dim, hid3 size is 10, and hid2 is geometric mean
            hid1_size = self.obs_dim * self.hid1_mult  # default multipler 10 chosen empirically on 'Hopper-v1'
            hid3_size = 5  # 5 chosen empirically on 'Hopper-v1'
            hid2_size = int(np.sqrt(hid1_size * hid3_size))
            # heuristic to set learning rate based on NN size (tuned on 'Hopper-v1')
            self.lr = 1e-2 / np.sqrt(hid2_size)  # 1e-3 empirically determined
            print('Value Params -- h1: {}, h2: {}, h3: {}, lr: {:.3g}'
                  .format(hid1_size, hid2_size, hid3_size, self.lr))
            # 3 hidden layers with tanh activations
            out = tf.layers.dense(self.obs_ph, hid1_size, tf.tanh,
                                  kernel_initializer=tf.random_normal_initializer(
                                      stddev=np.sqrt(1 / self.obs_dim)), name="h1")

            out = tf.layers.dense(out, hid2_size, tf.tanh,
                                  kernel_initializer=tf.random_normal_initializer(
                                      stddev=np.sqrt(1 / hid1_size)), name="h2")
            out = tf.layers.dense(out, hid3_size, tf.tanh,
                                  kernel_initializer=tf.random_normal_initializer(
                                      stddev=np.sqrt(1 / hid2_size)), name="h3")
            out = tf.layers.dense(out, 1,
                                  kernel_initializer=tf.random_normal_initializer(
                                      stddev=np.sqrt(1 / hid3_size)), name='output')
            self.out = tf.squeeze(out)
            self.loss = tf.reduce_mean(tf.square(self.out - self.val_ph))  # squared loss
            optimizer = tf.train.AdamOptimizer(self.lr)
            self.train_op = optimizer.minimize(self.loss)

            #self.init = tf.global_variables_initializer()
            self.h1_w, self.h1_b = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, self._scope_name+'/h1')
            self.h2_w, self.h2_b = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, self._scope_name+'/h2')
            self.h3_w, self.h3_b = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, self._scope_name+'/h3')
            self.output_w, self.output_b =tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, self._scope_name+'/output')

            scope.reuse_variables()
           
               
        #self.sess = tf.Session(graph=self.g)
        #self.sess.run(self.init)
    def fit_for_global(self, x, y, logger):
        """ Fit model to current data batch + previous data batch

        Args:
            x: features
            y: target
            logger: logger to save training loss and % explained variance
        """
        num_batches = max(x.shape[0] // 256, 1)
        batch_size = x.shape[0] // num_batches
        y_hat = self.predict(sess, x)  # check explained variance prior to update
        old_exp_var = 1 - np.var(y - y_hat)/np.var(y)
        if self.replay_buffer_x is None:
            x_train, y_train = x, y
        else:
            x_train = np.concatenate([x, self.replay_buffer_x])
            y_train = np.concatenate([y, self.replay_buffer_y])
        self.replay_buffer_x = x
        self.replay_buffer_y = y
        for e in range(self.epochs):
            x_train, y_train = shuffle(x_train, y_train)
            for j in range(num_batches):
                start = j * batch_size
                end = (j + 1) * batch_size
                feed_dict = {self.obs_ph: x_train[start:end, :],
                             self.val_ph: y_train[start:end]}
                _, l = sess.run([self.train_op, self.loss], feed_dict=feed_dict)
        y_hat = self.predict(sess, x)
        loss = np.mean(np.square(y_hat - y))         # explained variance after update
        exp_var = 1 - np.var(y - y_hat) / np.var(y)  # diagnose over-fitting of val func

        logger.log({'ValFuncLoss': loss,
                    'ExplainedVarNew': exp_var,
                    'ExplainedVarOld': old_exp_var})

    def fit(self, sess, x, y, logger):
        """ Fit model to current data batch + previous data batch

        Args:
            x: features
            y: target
            logger: logger to save training loss and % explained variance
        """
        num_batches = max(x.shape[0] // 256, 1)
        batch_size = x.shape[0] // num_batches
        y_hat = self.predict(sess, x)  # check explained variance prior to update
        old_exp_var = 1 - np.var(y - y_hat)/np.var(y)
        if self.replay_buffer_x is None:
            x_train, y_train = x, y
        else:
            x_train = np.concatenate([x, self.replay_buffer_x])
            y_train = np.concatenate([y, self.replay_buffer_y])
        self.replay_buffer_x = x
        self.replay_buffer_y = y
        for e in range(self.epochs):
            x_train, y_train = shuffle(x_train, y_train)
            for j in range(num_batches):
                start = j * batch_size
                end = (j + 1) * batch_size
                feed_dict = {self.obs_ph: x_train[start:end, :],
                             self.val_ph: y_train[start:end]}
                _, l = sess.run([self.train_op, self.loss], feed_dict=feed_dict)
        y_hat = self.predict(sess, x)
        loss = np.mean(np.square(y_hat - y))         # explained variance after update
        exp_var = 1 - np.var(y - y_hat) / np.var(y)  # diagnose over-fitting of val func

        logger.log({'ValFuncLoss': loss,
                    'ExplainedVarNew': exp_var,
                    'ExplainedVarOld': old_exp_var})

    def predict(self, sess, x):
        """ Predict method """
        feed_dict = {self.obs_ph: x}
        y_hat = sess.run(self.out, feed_dict=feed_dict)

        return np.squeeze(y_hat)

    #def close_sess(self):
    #    """ Close TensorFlow session """
    #    sess.close()
    def get_vars(self):
        return [self.h1_w, self.h1_b,
                self.h2_w, self.h2_b,
                self.h3_w, self.h3_b,
                self.output_w, self.output_b ]

#        weights = []

        #name = []
        #for tensor in self.g.as_graph_def().node:
        #    name.append(tensor.name)
        #print(name)

        #with self.g.as_default() as g:
#        with tf.variable_scope(self._scope_name) as scope:
#            weights.extend(tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope))
            # weights.append(g.get_tensor_by_name('h1/kernel:0'))
            # weights.append(g.get_tensor_by_name('h1/bias:0'))
            # weights.append(g.get_tensor_by_name('h2/kernel:0'))
            # weights.append(g.get_tensor_by_name('h2/bias:0'))
            # weights.append(g.get_tensor_by_name('h3/kernel:0'))
            # weights.append(g.get_tensor_by_name('h3/bias:0'))
  


#        return weights

    def sync_from(self, shared_nn, name=None):
        if shared_nn != None:
            src_vars = shared_nn.get_vars()
            dst_vars = self.get_vars()
            sync_ops = []

            with tf.name_scope(name, self._scope_name, []) as name:
        
                for(src_var, dst_var) in zip(src_vars, dst_vars):
  
                    sync_op = tf.assign(dst_var, src_var)
                    sync_ops.append(sync_op)
        
            return tf.group(*sync_ops, name=name)
        else:
            return None



[((35, 25, 39, 46), 'tensorflow.gradients', 'tf.gradients', (), '', True, 'import tensorflow as tf\n'), ((163, 15, 163, 32), 'numpy.squeeze', 'np.squeeze', ({(163, 26, 163, 31): 'y_hat'}, {}), '(y_hat)', True, 'import numpy as np\n'), ((47, 13, 47, 48), 'tensorflow.variable_scope', 'tf.variable_scope', ({(47, 31, 47, 47): 'self._scope_name'}, {}), '(self._scope_name)', True, 'import tensorflow as tf\n'), ((48, 26, 48, 89), 'tensorflow.placeholder', 'tf.placeholder', ({(48, 41, 48, 51): 'tf.float32', (48, 53, 48, 73): '(None, self.obs_dim)', (48, 75, 48, 88): '"""obs_valfunc"""'}, {}), "(tf.float32, (None, self.obs_dim), 'obs_valfunc')", True, 'import tensorflow as tf\n'), ((49, 26, 49, 76), 'tensorflow.placeholder', 'tf.placeholder', ({(49, 41, 49, 51): 'tf.float32', (49, 53, 49, 60): '(None,)', (49, 62, 49, 75): '"""val_valfunc"""'}, {}), "(tf.float32, (None,), 'val_valfunc')", True, 'import tensorflow as tf\n'), ((72, 23, 72, 38), 'tensorflow.squeeze', 'tf.squeeze', ({(72, 34, 72, 37): 'out'}, {}), '(out)', True, 'import tensorflow as tf\n'), ((74, 24, 74, 55), 'tensorflow.train.AdamOptimizer', 'tf.train.AdamOptimizer', ({(74, 47, 74, 54): 'self.lr'}, {}), '(self.lr)', True, 'import tensorflow as tf\n'), ((78, 35, 78, 110), 'tensorflow.get_collection', 'tf.get_collection', ({(78, 53, 78, 85): 'tf.GraphKeys.TRAINABLE_VARIABLES', (78, 87, 78, 109): "self._scope_name + '/h1'"}, {}), "(tf.GraphKeys.TRAINABLE_VARIABLES, self._scope_name + '/h1')", True, 'import tensorflow as tf\n'), ((79, 35, 79, 110), 'tensorflow.get_collection', 'tf.get_collection', ({(79, 53, 79, 85): 'tf.GraphKeys.TRAINABLE_VARIABLES', (79, 87, 79, 109): "self._scope_name + '/h2'"}, {}), "(tf.GraphKeys.TRAINABLE_VARIABLES, self._scope_name + '/h2')", True, 'import tensorflow as tf\n'), ((80, 35, 80, 110), 'tensorflow.get_collection', 'tf.get_collection', ({(80, 53, 80, 85): 'tf.GraphKeys.TRAINABLE_VARIABLES', (80, 87, 80, 109): "self._scope_name + '/h3'"}, {}), "(tf.GraphKeys.TRAINABLE_VARIABLES, self._scope_name + '/h3')", True, 'import tensorflow as tf\n'), ((81, 42, 81, 121), 'tensorflow.get_collection', 'tf.get_collection', ({(81, 60, 81, 92): 'tf.GraphKeys.TRAINABLE_VARIABLES', (81, 94, 81, 120): "self._scope_name + '/output'"}, {}), "(tf.GraphKeys.TRAINABLE_VARIABLES, self._scope_name +\n    '/output')", True, 'import tensorflow as tf\n'), ((103, 22, 103, 63), 'numpy.concatenate', 'np.concatenate', ({(103, 37, 103, 62): '[x, self.replay_buffer_x]'}, {}), '([x, self.replay_buffer_x])', True, 'import numpy as np\n'), ((104, 22, 104, 63), 'numpy.concatenate', 'np.concatenate', ({(104, 37, 104, 62): '[y, self.replay_buffer_y]'}, {}), '([y, self.replay_buffer_y])', True, 'import numpy as np\n'), ((108, 31, 108, 56), 'sklearn.utils.shuffle', 'shuffle', ({(108, 39, 108, 46): 'x_train', (108, 48, 108, 55): 'y_train'}, {}), '(x_train, y_train)', False, 'from sklearn.utils import shuffle\n'), ((116, 23, 116, 43), 'numpy.square', 'np.square', ({(116, 33, 116, 42): 'y_hat - y'}, {}), '(y_hat - y)', True, 'import numpy as np\n'), ((138, 22, 138, 63), 'numpy.concatenate', 'np.concatenate', ({(138, 37, 138, 62): '[x, self.replay_buffer_x]'}, {}), '([x, self.replay_buffer_x])', True, 'import numpy as np\n'), ((139, 22, 139, 63), 'numpy.concatenate', 'np.concatenate', ({(139, 37, 139, 62): '[y, self.replay_buffer_y]'}, {}), '([y, self.replay_buffer_y])', True, 'import numpy as np\n'), ((143, 31, 143, 56), 'sklearn.utils.shuffle', 'shuffle', ({(143, 39, 143, 46): 'x_train', (143, 48, 143, 55): 'y_train'}, {}), '(x_train, y_train)', False, 'from sklearn.utils import shuffle\n'), ((151, 23, 151, 43), 'numpy.square', 'np.square', ({(151, 33, 151, 42): 'y_hat - y'}, {}), '(y_hat - y)', True, 'import numpy as np\n'), ((208, 19, 208, 49), 'tensorflow.group', 'tf.group', (), '', True, 'import tensorflow as tf\n'), ((53, 28, 53, 58), 'numpy.sqrt', 'np.sqrt', ({(53, 36, 53, 57): 'hid1_size * hid3_size'}, {}), '(hid1_size * hid3_size)', True, 'import numpy as np\n'), ((55, 29, 55, 47), 'numpy.sqrt', 'np.sqrt', ({(55, 37, 55, 46): 'hid2_size'}, {}), '(hid2_size)', True, 'import numpy as np\n'), ((73, 39, 73, 72), 'tensorflow.square', 'tf.square', ({(73, 49, 73, 71): 'self.out - self.val_ph'}, {}), '(self.out - self.val_ph)', True, 'import tensorflow as tf\n'), ((99, 26, 99, 43), 'numpy.var', 'np.var', ({(99, 33, 99, 42): '(y - y_hat)'}, {}), '(y - y_hat)', True, 'import numpy as np\n'), ((99, 44, 99, 53), 'numpy.var', 'np.var', ({(99, 51, 99, 52): 'y'}, {}), '(y)', True, 'import numpy as np\n'), ((117, 22, 117, 39), 'numpy.var', 'np.var', ({(117, 29, 117, 38): '(y - y_hat)'}, {}), '(y - y_hat)', True, 'import numpy as np\n'), ((117, 42, 117, 51), 'numpy.var', 'np.var', ({(117, 49, 117, 50): 'y'}, {}), '(y)', True, 'import numpy as np\n'), ((134, 26, 134, 43), 'numpy.var', 'np.var', ({(134, 33, 134, 42): '(y - y_hat)'}, {}), '(y - y_hat)', True, 'import numpy as np\n'), ((134, 44, 134, 53), 'numpy.var', 'np.var', ({(134, 51, 134, 52): 'y'}, {}), '(y)', True, 'import numpy as np\n'), ((152, 22, 152, 39), 'numpy.var', 'np.var', ({(152, 29, 152, 38): '(y - y_hat)'}, {}), '(y - y_hat)', True, 'import numpy as np\n'), ((152, 42, 152, 51), 'numpy.var', 'np.var', ({(152, 49, 152, 50): 'y'}, {}), '(y)', True, 'import numpy as np\n'), ((201, 17, 201, 58), 'tensorflow.name_scope', 'tf.name_scope', ({(201, 31, 201, 35): 'name', (201, 37, 201, 53): 'self._scope_name', (201, 55, 201, 57): '[]'}, {}), '(name, self._scope_name, [])', True, 'import tensorflow as tf\n'), ((205, 30, 205, 57), 'tensorflow.assign', 'tf.assign', ({(205, 40, 205, 47): 'dst_var', (205, 49, 205, 56): 'src_var'}, {}), '(dst_var, src_var)', True, 'import tensorflow as tf\n'), ((61, 45, 61, 70), 'numpy.sqrt', 'np.sqrt', ({(61, 53, 61, 69): '1 / self.obs_dim'}, {}), '(1 / self.obs_dim)', True, 'import numpy as np\n'), ((65, 45, 65, 67), 'numpy.sqrt', 'np.sqrt', ({(65, 53, 65, 66): '1 / hid1_size'}, {}), '(1 / hid1_size)', True, 'import numpy as np\n'), ((68, 45, 68, 67), 'numpy.sqrt', 'np.sqrt', ({(68, 53, 68, 66): '1 / hid2_size'}, {}), '(1 / hid2_size)', True, 'import numpy as np\n'), ((71, 45, 71, 67), 'numpy.sqrt', 'np.sqrt', ({(71, 53, 71, 66): '1 / hid3_size'}, {}), '(1 / hid3_size)', True, 'import numpy as np\n')]