# Path: agents/capacities.py
# def swish(x):
# def eps_greedy(inputs_t, q_preds_t, nb_actions, N0, min_eps, nb_state=None):
# def tabular_eps_greedy(inputs_t, q_preds_t, nb_states, nb_actions, N0, min_eps):
# def tabular_UCB(Qs_t, inputs_t):
# def eligibility_traces(Qs_t, states_t, actions_t, discount, lambda_value):
# def eligibility_dutch_traces(Qs_t, states_t, actions_t, lr, discount, lambda_value):
# def get_expected_rewards(episode_rewards, discount=.99):
# def tf_get_n_step_expected_rewards(episode_rewards_t, estimates_t, discount, n_step):
# def get_n_step_expected_rewards(episode_rewards, estimates, discount=.99, n_step=1):
# def get_n_step_expected_rewards_mat(episode_rewards, estimates, discount=.99, n_step=1):
# def get_lambda_expected_rewards(episode_rewards, estimates, discount=.99, lambda_value=.9):
# def get_mc_target(rewards_t, discount):
# def get_td_target(Qs_t, rewards_t, next_states_t, next_actions_t, discount):
# def get_td_n_target(Qs_t, rewards_t, next_state_t, next_action_t, discount, n_step):
# def get_q_learning_target(Qs_t, rewards_t, next_states_t, discount):
# def get_expected_sarsa_target(Qs_t, rewards_t, next_states_t, next_probs_t, discount):
# def get_sigma_target(Qs_t, sigma, rewards_t, next_states_t, next_actions_t, next_probs_t, discount):
# def tabular_learning(Qs_t, states_t, actions_t, targets):
# def tabular_learning_with_lr(init_lr, decay_steps, Qs_t, states_t, actions_t, targets):
# def counter(name):
# def fix_scope(from_scope):
# def build_batches(dtKeys, sequence_history, batch_size):
# def policy(network_params, inputs, summary_collections=None):
# def value_f(network_params, inputs, summary_collections=None):
# def predictive_model(network_params, m_inputs, dynamic_batch_size, m_initial_state=None, summary_collections=None):
# def projection(proj_params, inputs, summary_collections=None):
#         N = tf.Variable(1., trainable=False, dtype=tf.float32, name='N')
#         N = tf.Variable(tf.ones(shape=[nb_state]), name='N', trainable=False)
#     W1 = tf.get_variable('W1'
#         , shape=[ network_params['nb_inputs'], network_params['nb_units'] ]
#         , initializer=tf.random_normal_initializer(mean=network_params['initial_mean'], stddev=network_params['initial_stddev'])
#     )
#     W2 = tf.get_variable('W2'
#         , shape=[ network_params['nb_units'], network_params['nb_units'] ]
#         , initializer=tf.random_normal_initializer(mean=network_params['initial_mean'], stddev=network_params['initial_stddev'])
#     )
#     W3 = tf.get_variable('W3'
#         , shape=[ network_params['nb_units'], network_params['nb_outputs'] ]
#         , initializer=tf.random_normal_initializer(mean=network_params['initial_mean'], stddev=network_params['initial_stddev'])
#     )
#     W1 = tf.get_variable('W1'
#         , shape=[ network_params['nb_inputs'], network_params['nb_units'] ]
#         , initializer=w_init
#     )
#     W2 = tf.get_variable('W2'
#         , shape=[ network_params['nb_units'], network_params['nb_units'] ]
#         , initializer=w_init
#     )
#     W3 = tf.get_variable('W3'
#         , shape=[ network_params['nb_units'], network_params['nb_outputs'] ]
#         , initializer=w_init
#     )
#     W1 = tf.get_variable('W1'
#         , shape=[ proj_params['nb_units'], proj_params['nb_units'] ]
#         , initializer=w_init
#     )
#     W2 = tf.get_variable('W2'
#         , shape=[ proj_params['nb_units'], proj_params['nb_actions'] ]
#         , initializer=w_init
#     )


import os, sys, unittest, timeit
import numpy as np
import tensorflow as tf
from agents import capacities

            et, update_et_op, reset_et_op = capacities.eligibility_traces(Qs_t, states_t, actions_t, discount, lambda_value)

            with tf.Session() as sess:
                sess.run(tf.global_variables_initializer())

                _ = sess.run([update_et_op], feed_dict={
                    states_t: [0],
                    actions_t: [1]
                })
                self.assertEqual(np.array_equal(sess.run(et), [[ 0. , 1.], [ 0. , 0.], [ 0. , 0.]]), True)
                _ = sess.run([reset_et_op])
                self.assertEqual(np.array_equal(sess.run(et), [[ 0. , 0.], [ 0. , 0.], [ 0. , 0.]]), True)

    def test_get_mc_target(self):
        discount = .5
        rewards = [1, 1, 2]

        rewards_plh = tf.placeholder(tf.float32, shape=[None])
        target_t = capacities.get_mc_target(rewards_plh, discount)
        with tf.Session() as sess:
            target = sess.run(target_t, feed_dict={rewards_plh: rewards})
            # print('10000 call:', timeit.timeit(lambda: sess.run(target_t, feed_dict={rewards_plh: rewards}), number=10000))

        self.assertEqual(np.sum(np.isclose(target, [2, 2, 2])) == 3, True)

    def test_get_expected_rewards_with_discount(self):
        rewards = [1, 1, 2]
        discount = .5
        

# Path: agents/capacities.py
# def swish(x):
# def eps_greedy(inputs_t, q_preds_t, nb_actions, N0, min_eps, nb_state=None):
# def tabular_eps_greedy(inputs_t, q_preds_t, nb_states, nb_actions, N0, min_eps):
# def tabular_UCB(Qs_t, inputs_t):
# def eligibility_traces(Qs_t, states_t, actions_t, discount, lambda_value):
# def eligibility_dutch_traces(Qs_t, states_t, actions_t, lr, discount, lambda_value):
# def get_expected_rewards(episode_rewards, discount=.99):
# def tf_get_n_step_expected_rewards(episode_rewards_t, estimates_t, discount, n_step):
# def get_n_step_expected_rewards(episode_rewards, estimates, discount=.99, n_step=1):
# def get_n_step_expected_rewards_mat(episode_rewards, estimates, discount=.99, n_step=1):
# def get_lambda_expected_rewards(episode_rewards, estimates, discount=.99, lambda_value=.9):
# def get_mc_target(rewards_t, discount):
# def get_td_target(Qs_t, rewards_t, next_states_t, next_actions_t, discount):
# def get_td_n_target(Qs_t, rewards_t, next_state_t, next_action_t, discount, n_step):
# def get_q_learning_target(Qs_t, rewards_t, next_states_t, discount):
# def get_expected_sarsa_target(Qs_t, rewards_t, next_states_t, next_probs_t, discount):
# def get_sigma_target(Qs_t, sigma, rewards_t, next_states_t, next_actions_t, next_probs_t, discount):
# def tabular_learning(Qs_t, states_t, actions_t, targets):
# def tabular_learning_with_lr(init_lr, decay_steps, Qs_t, states_t, actions_t, targets):
# def counter(name):
# def fix_scope(from_scope):
# def build_batches(dtKeys, sequence_history, batch_size):
# def policy(network_params, inputs, summary_collections=None):
# def value_f(network_params, inputs, summary_collections=None):
# def predictive_model(network_params, m_inputs, dynamic_batch_size, m_initial_state=None, summary_collections=None):
# def projection(proj_params, inputs, summary_collections=None):
#         N = tf.Variable(1., trainable=False, dtype=tf.float32, name='N')
#         N = tf.Variable(tf.ones(shape=[nb_state]), name='N', trainable=False)
#     W1 = tf.get_variable('W1'
#         , shape=[ network_params['nb_inputs'], network_params['nb_units'] ]
#         , initializer=tf.random_normal_initializer(mean=network_params['initial_mean'], stddev=network_params['initial_stddev'])
#     )
#     W2 = tf.get_variable('W2'
#         , shape=[ network_params['nb_units'], network_params['nb_units'] ]
#         , initializer=tf.random_normal_initializer(mean=network_params['initial_mean'], stddev=network_params['initial_stddev'])
#     )
#     W3 = tf.get_variable('W3'
#         , shape=[ network_params['nb_units'], network_params['nb_outputs'] ]
#         , initializer=tf.random_normal_initializer(mean=network_params['initial_mean'], stddev=network_params['initial_stddev'])
#     )
#     W1 = tf.get_variable('W1'
#         , shape=[ network_params['nb_inputs'], network_params['nb_units'] ]
#         , initializer=w_init
#     )
#     W2 = tf.get_variable('W2'
#         , shape=[ network_params['nb_units'], network_params['nb_units'] ]
#         , initializer=w_init
#     )
#     W3 = tf.get_variable('W3'
#         , shape=[ network_params['nb_units'], network_params['nb_outputs'] ]
#         , initializer=w_init
#     )
#     W1 = tf.get_variable('W1'
#         , shape=[ proj_params['nb_units'], proj_params['nb_units'] ]
#         , initializer=w_init
#     )
#     W2 = tf.get_variable('W2'
#         , shape=[ proj_params['nb_units'], proj_params['nb_actions'] ]
#         , initializer=w_init
#     )


# Path: tests/agents/test_capacities.py
import os, sys, unittest, timeit
import numpy as np
import tensorflow as tf
from agents import capacities


            et, update_et_op, reset_et_op = capacities.eligibility_traces(Qs_t, states_t, actions_t, discount, lambda_value)

            with tf.Session() as sess:
                sess.run(tf.global_variables_initializer())

                _ = sess.run([update_et_op], feed_dict={
                    states_t: [0],
                    actions_t: [1]
                })
                self.assertEqual(np.array_equal(sess.run(et), [[ 0. , 1.], [ 0. , 0.], [ 0. , 0.]]), True)
                _ = sess.run([reset_et_op])
                self.assertEqual(np.array_equal(sess.run(et), [[ 0. , 0.], [ 0. , 0.], [ 0. , 0.]]), True)

    def test_get_mc_target(self):
        discount = .5
        rewards = [1, 1, 2]

        rewards_plh = tf.placeholder(tf.float32, shape=[None])
        target_t = capacities.get_mc_target(rewards_plh, discount)
        with tf.Session() as sess:
            target = sess.run(target_t, feed_dict={rewards_plh: rewards})
            # print('10000 call:', timeit.timeit(lambda: sess.run(target_t, feed_dict={rewards_plh: rewards}), number=10000))

        self.assertEqual(np.sum(np.isclose(target, [2, 2, 2])) == 3, True)

    def test_get_expected_rewards_with_discount(self):
        rewards = [1, 1, 2]
        discount = .5
        

        expected_rewards = capacities.get_expected_rewards(rewards, discount)