# Path: util/logger.py
# class _MyFormatter(logging.Formatter):
#     def format(self, record):
# def set_file_handler(path=None, prefix='', time_str=''):
# _LOGGING_METHOD = ['info', 'warning', 'error', 'critical',
#                    'warn', 'exception', 'debug']
# 
# Path: util/utils.py
# def gauss_selfKL_firstfixed(mu, logstd):
# def gauss_log_prob(mu, logstd, x):
# def gauss_KL(mu1, logstd1, mu2, logstd2):
# def gauss_ent(mu, logstd):
# def discount(x, gamma):
# def cat_sample(prob_nk, seed=1234):
# def slice_2d(x, inds0, inds1):
# def var_shape(x):
# def numel(x):
# def flatgrad(loss, var_list):
# def conjugate_gradient(f_Ax, b, cg_iters=10, residual_tol=1e-10):
# def linesearch(f, x, fullstep, expected_improve_rate):
#   def __init__(self, session, var_list):
#   def __call__(self, theta):
#   def __init__(self, session, var_list):
#   def __call__(self):
#   def __init__(self, session, var_list):
#   def __call__(self):
#   def __init__(self, session, var_list):
#   def __call__(self, weights):
# def xavier_initializer(self, shape):
# def fully_connected(input_layer, input_size, output_size, weight_init,
#                     bias_init, scope, trainable):
#   N = prob_nk.shape[0]
# class SetFromFlat(object):
# class GetFlat(object):
# class GetPolicyWeights(object):
# class SetPolicyWeights(object):
# 
# Path: graph_util/nervenetplus_util.py
# def nervenetplus_step_assign(rollout_data, nstep):
# 
# Path: graph_util/graph_data_util.py
# _ABS_BASE_PATH = init_path.get_abs_base_dir()
# def construct_graph_input_feeddict(node_info,
#                                    obs_n,
#                                    receive_idx,
#                                    send_idx,
#                                    node_type_idx,
#                                    inverse_node_type_idx,
#                                    output_type_idx,
#                                    inverse_output_type_idx,
#                                    last_batch_size,
#                                    request_data=['ob', 'idx']):
# def _construct_index(node_info, batch_size):
# def _get_obs(node_info, raw_ob, batch_size):
# def _add_offset(idx, offset, batch_size, if_reshape):
# def _add_inverse_offset(extype_offset, intype_offset, self_offset, batch_size):


import init_path
import numpy as np
from util import logger
from util import utils
from graph_util import nervenetplus_util
from graph_util import graph_data_util
    }, nervenetplus_batch_pos


def generate_advantage(data_dict, baseline_network):
    '''
        @brief: calculate the parameters for the advantage function
    '''

    for path in data_dict:
        # the predicted value function (baseline function)
        path["baseline"] = baseline_network.predict(path)

    advantage_method = baseline_network.args.advantage_method
    gamma = baseline_network.args.gamma
    gae_lam = baseline_network.args.gae_lam

    # esitmate the advantages
    if advantage_method == 'raw':
        for path in data_dict:
            # the gamma discounted rollout value function
            path["returns"] = utils.discount(path["rewards"], gamma)
            path["advantage"] = path["returns"] - path["baseline"]
            path['target_return'] = path['returns']
    else:
        assert advantage_method == 'gae', logger.error(
            'invalid advantage estimation method: {}'.format(advantage_method)
        )

        for path in data_dict:
            # the gamma discounted rollout value function

# Path: util/logger.py
# class _MyFormatter(logging.Formatter):
#     def format(self, record):
# def set_file_handler(path=None, prefix='', time_str=''):
# _LOGGING_METHOD = ['info', 'warning', 'error', 'critical',
#                    'warn', 'exception', 'debug']
# 
# Path: util/utils.py
# def gauss_selfKL_firstfixed(mu, logstd):
# def gauss_log_prob(mu, logstd, x):
# def gauss_KL(mu1, logstd1, mu2, logstd2):
# def gauss_ent(mu, logstd):
# def discount(x, gamma):
# def cat_sample(prob_nk, seed=1234):
# def slice_2d(x, inds0, inds1):
# def var_shape(x):
# def numel(x):
# def flatgrad(loss, var_list):
# def conjugate_gradient(f_Ax, b, cg_iters=10, residual_tol=1e-10):
# def linesearch(f, x, fullstep, expected_improve_rate):
#   def __init__(self, session, var_list):
#   def __call__(self, theta):
#   def __init__(self, session, var_list):
#   def __call__(self):
#   def __init__(self, session, var_list):
#   def __call__(self):
#   def __init__(self, session, var_list):
#   def __call__(self, weights):
# def xavier_initializer(self, shape):
# def fully_connected(input_layer, input_size, output_size, weight_init,
#                     bias_init, scope, trainable):
#   N = prob_nk.shape[0]
# class SetFromFlat(object):
# class GetFlat(object):
# class GetPolicyWeights(object):
# class SetPolicyWeights(object):
# 
# Path: graph_util/nervenetplus_util.py
# def nervenetplus_step_assign(rollout_data, nstep):
# 
# Path: graph_util/graph_data_util.py
# _ABS_BASE_PATH = init_path.get_abs_base_dir()
# def construct_graph_input_feeddict(node_info,
#                                    obs_n,
#                                    receive_idx,
#                                    send_idx,
#                                    node_type_idx,
#                                    inverse_node_type_idx,
#                                    output_type_idx,
#                                    inverse_output_type_idx,
#                                    last_batch_size,
#                                    request_data=['ob', 'idx']):
# def _construct_index(node_info, batch_size):
# def _get_obs(node_info, raw_ob, batch_size):
# def _add_offset(idx, offset, batch_size, if_reshape):
# def _add_inverse_offset(extype_offset, intype_offset, self_offset, batch_size):


# Path: graph_util/gnn_io_util.py
import init_path
import numpy as np
from util import logger
from util import utils
from graph_util import nervenetplus_util
from graph_util import graph_data_util

    }, nervenetplus_batch_pos


def generate_advantage(data_dict, baseline_network):
    '''
        @brief: calculate the parameters for the advantage function
    '''

    for path in data_dict:
        # the predicted value function (baseline function)
        path["baseline"] = baseline_network.predict(path)

    advantage_method = baseline_network.args.advantage_method
    gamma = baseline_network.args.gamma
    gae_lam = baseline_network.args.gae_lam

    # esitmate the advantages
    if advantage_method == 'raw':
        for path in data_dict:
            # the gamma discounted rollout value function
            path["returns"] = utils.discount(path["rewards"], gamma)
            path["advantage"] = path["returns"] - path["baseline"]
            path['target_return'] = path['returns']
    else:
        assert advantage_method == 'gae', logger.error(
            'invalid advantage estimation method: {}'.format(advantage_method)
        )

        for path in data_dict:
            # the gamma discounted rollout value function

            path["returns"] = utils.discount(path["rewards"], gamma)