import numpy as np
import tensorflow as tf
import tflearn
import tensorflow.contrib.layers as tl

#for each data, we only consider at most 2 candidate
K = 2 

#??
GAMMA = 0.99

#action space dimension: equivalent to the number of candidates plus 1 for no partition
A_DIM = K + 1

#??
ENTROPY_WEIGHT = 0.1

#??
ENTROPY_EPS = 1e-6

#state space dimension
S_DIM = 4 * K + 7


class ActorNetwork(object):

#Input: state, including:
#  ---K candidate distances: the level of indirectness from the candidate lambda to current computation
#  ---K candidate frequency
#  ---K candidate relative selectivity: the historical average of the candidate lambda's output size to current input size
#  ---Current input size
#  ---Number of nodes in the cluster
#Output: probability distribution of all actions

    def __init__(self, sess, state_dim, action_dim, learning_rate):
        self.sess = sess
        self.s_dim = state_dim
        self.a_dim = action_dim
        self.lr_rate = learning_rate

        #create the actor network
        self.inputs = tf.placeholder(tf.float32, [None, self.s_dim])
        self.out = self.create_actor_network()

        #get all network parameters
        self.network_params = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope='actor')

        #set all network parameters
        self.input_network_params = []
        for param in self.network_params:
            self.input_network_params.append(
                tf.placeholder(tf.float32, shape=param.get_shape()))
        self.set_network_params_op = []
        for idx, param in enumerate(self.input_network_params):
            self.set_network_params_op.append(self.network_params[idx].assign(param))

        #selected action, 0-1 vector
        self.acts = tf.placeholder(tf.float32, [None, self.a_dim])

        #this gradient will be provided by the critic network
        self.act_grad_weights = tf.placeholder(tf.float32, [None, 1])

        print ("self.act_grad_weights: ")
        print self.act_grad_weights

        print ("self.act_grad_weights.shape: ")
        print self.act_grad_weights.shape


        #compute the objective 
        self.obj = tf.reduce_sum(tf.multiply(
                       tf.log(tf.reduce_sum(tf.multiply(self.out, self.acts),
                                            reduction_indices=1, keep_dims=True)),
                       -self.act_grad_weights)) \
                   + ENTROPY_WEIGHT * tf.reduce_sum(tf.multiply(self.out,
                                                           tf.log(self.out + ENTROPY_EPS)))

        #combine the gradients here
        self.actor_gradients = tf.gradients(self.obj, self.network_params)

        #optimization op
        self.optimize = tf.train.RMSPropOptimizer(self.lr_rate).\
            apply_gradients(zip(self.actor_gradients, self.network_params))

    def create_actor_network(self):
        with tf.variable_scope('actor'):
            print ("inputs:")
            print (self.inputs)
            hid_1 = tl.fully_connected(self.inputs, 16, activation_fn=tf.nn.leaky_relu)
            hid_2 = tl.fully_connected(hid_1, 8, activation_fn=tf.nn.leaky_relu)
            out = tl.fully_connected(hid_2, self.a_dim, activation_fn=tf.nn.softmax)
            print ("out:")
            print (out)
            return out

    def train(self, inputs, acts, act_grad_weights):

        self.sess.run(self.optimize, feed_dict={
            self.inputs: inputs,
            self.acts: acts,
            self.act_grad_weights: act_grad_weights
        })

    def predict(self, inputs):
        print("predict inputs in actor network:")
        print(inputs)
        return self.sess.run(self.out, feed_dict={
            self.inputs: inputs
        })

    def get_gradients(self, inputs, acts, act_grad_weights):
        print ("input: ")
        print (inputs)
       
        print ("acts: ")
        print (acts)

        print ("act_grad_weights: ")
        print (act_grad_weights)

        return self.sess.run(self.actor_gradients, feed_dict={
            self.inputs: inputs,
            self.acts: acts,
            self.act_grad_weights: act_grad_weights
        })

    def apply_gradients(self, actor_gradients):
        return self.sess.run(self.optimize, feed_dict={
            i: d for i, d in zip(self.actor_gradients, actor_gradients)
        })

    def get_network_params(self):
        return self.sess.run(self.network_params)

    def set_network_params(self, input_network_params):
        self.sess.run(self.set_network_params_op, feed_dict={
            i: d for i, d in zip(self.input_network_params, input_network_params) 
        })


class CriticNetwork(object):

    #input: the state and action 
    #output: value function V(s), which represents the cumulative reward in a long run

    def __init__(self, sess, state_dim, learning_rate):
        self.sess = sess
        self.s_dim = state_dim
        self.lr_rate = learning_rate

        #create the critic network
        self.inputs = tf.placeholder(tf.float32, [None, self.s_dim])
        self.out = self.create_critic_network()

        #get all network parameters
        self.network_params = \
            tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope='critic')

        #set all network parameters
        self.input_network_params = []
        for param in self.network_params:
            self.input_network_params.append(
                tf.placeholder(tf.float32, shape=param.get_shape()))
        self.set_network_params_op = []
        for idx, param in enumerate(self.input_network_params):
            self.set_network_params_op.append(self.network_params[idx].assign(param))

        #network target V(s)
        self.td_target = tf.placeholder(tf.float32, [None, 1])
          
        #temporal difference, will also be weights for actor_gradients
        self.td = tf.subtract(self.td_target, self.out)

        #mean square error
        #self.loss = tf.reduce_mean(tf.square(self.td_target - self.out))
        self.loss = tflearn.mean_square(self.td_target, self.out)
        #compute critic gradient
        self.critic_gradients = tf.gradients(self.loss, self.network_params)

        #optimize op
        self.optimize = tf.train.RMSPropOptimizer(self.lr_rate).\
            apply_gradients(zip(self.critic_gradients, self.network_params))

    def create_critic_network(self):
        with tf.variable_scope('critic'):
            
            hid_1 = tflearn.fully_connected(self.inputs, 16, activation='LeakyReLU')
            hid_2 = tflearn.fully_connected(hid_1, 8, activation='LeakyReLU')
            out = tflearn.fully_connected(hid_2, 1, activation='linear')
            return out

    def train(self, inputs, td_target):
        return self.sess.run([self.loss, self.optimize], feed_dict={
            self.inputs: inputs,
            self.td_target: td_target
        })

    def predict(self, inputs):
        print ("inputs in critic network:")
        print (inputs)
        return self.sess.run(self.out, feed_dict={
            self.inputs: inputs
        })

    def get_td(self, inputs, td_target):
        return self.sess.run(self.td, feed_dict={
            self.inputs: inputs,
            self.td_target: td_target
        })

    def get_gradients(self, inputs, td_target):
        return self.sess.run(self.critic_gradients, feed_dict={
            self.inputs: inputs,
            self.td_target: td_target
        })

    def apply_gradients(self, critic_gradients):
        return self.sess.run(self.optimize, feed_dict={
            i: d for i, d in zip(self.critic_gradients, critic_gradients)
        })

    def get_network_params(self):
        return self.sess.run(self.network_params)

    def set_network_params(self, input_network_params):
        self.sess.run(self.set_network_params_op, feed_dict={
            i: d for i, d in zip(self.input_network_params, input_network_params)
        })


def compute_gradients(s_batch, a_batch, r_batch, terminal, actor, critic):
   
    #batch of s, a, r is from samples in a sequence
    #the format is in np.array([batch_size, s/a/r_dim])
    #terminal is True when sequence ends as a terminal state
    print s_batch.shape[0]
    print a_batch.shape[0]
    print r_batch.shape[0]

    assert s_batch.shape[0] == a_batch.shape[0]
    assert s_batch.shape[0] == r_batch.shape[0]
    ba_size = s_batch.shape[0]

    v_batch = critic.predict(s_batch)
    R_batch = np.zeros(r_batch.shape)

    if terminal:
        R_batch[-1, 0] = 0  # terminal state
    else:
        R_batch[-1, 0] = v_batch[-1, 0]  # boot strap from last state

    for t in reversed(xrange(ba_size - 1)):
         #index = t+3
         #if index > 94:
         #   index = t
         #R_batch[t, 0] = r_batch[index] + GAMMA * R_batch[t + 1, 0]#8 is hardcoded for TPCH
         R_batch[t, 0] = r_batch[t] + GAMMA * R_batch[t + 1, 0]

    td_batch = R_batch - v_batch

    print ("R_batch: ")
    print R_batch

    print ("v_batch: ")
    print v_batch

    print ("td_batch: ")
    print td_batch


    print ("s_batch: ")
    print s_batch

    print ("a_batch: ")
    print a_batch

   

    actor_gradients = actor.get_gradients(s_batch, a_batch, td_batch)
    critic_gradients = critic.get_gradients(s_batch, R_batch)

    return actor_gradients, critic_gradients, td_batch


def discount(x, gamma):
    
    #Given vector x, computes a vector y such that
    #y[i] = x[i] + gamma * x[i+1] + gamma^2 x[i+2] + ...

    out = np.zeros(len(x))
    out[-1] = x[-1]
    for i in reversed(xrange(len(x)-1)):
        out[i] = x[i] + gamma*out[i+1]
    assert x.ndim >= 1
    # More efficient version:
    # scipy.signal.lfilter([1],[1,-gamma],x[::-1], axis=0)[::-1]
    return out

def compute_entropy(x):
    
    #Given vector x, computes the entropy
    #H(x) = - sum( p * log(p))
   
    H = 0.0
    for i in xrange(len(x)):
        if 0 < x[i] < 1:
            H -= x[i] * np.log(x[i])
    return H


def build_summaries():
    td_loss = tf.Variable(0.)
    tf.summary.scalar("TD_loss", td_loss)
    eps_total_reward = tf.Variable(0.)
    tf.summary.scalar("Eps_total_reward", eps_total_reward)

    summary_vars = [td_loss, eps_total_reward]
    summary_ops = tf.summary.merge_all()

    return summary_ops, summary_vars




