import tensorflow as tf

class ValueEstimator_Pendulum():
    def __init__(self, learning_rate=0.1, par_idx=0,scope="value_estimator"):
        w_init = tf.random_normal_initializer(0.,.1);
        with tf.variable_scope(scope+"_"+str(par_idx)):
            # state and target
            self.state = tf.placeholder(tf.float32, [None,num_state], "state")
            self.target = tf.placeholder(tf.float32, [None,1], name="target")

            # layers
            l_c = tf.layers.dense(self.state, 100, tf.nn.relu6, kernel_initializer=w_init, name='lc')
            self.value_estimate = tf.layers.dense(l_c, 1, kernel_initializer=w_init, name='v')  # estimated value for state
        
            # loss and optimizer
            self.loss = tf.reduce_mean(tf.square(tf.subtract(self.value_estimate, self.target)))
            self.optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate)
            self.train_op = self.optimizer.minimize(
                self.loss, global_step=tf.contrib.framework.get_global_step())        
    
    def predict(self, state, sess=None):
        sess = sess or tf.get_default_session()
        return sess.run(self.value_estimate, { self.state: [state] })[0][0]

    def update(self, state, target, sess=None):
        sess = sess or tf.get_default_session()
        feed_dict = { self.state: state, self.target: target }
        _, loss = sess.run([self.train_op, self.loss], feed_dict)
        return loss

