"""Deep Q learning graph

The functions in this file can are used to create the following functions:

======= act ========

    Function to chose an action given an observation

    Parameters
    ----------
    observation: object
        Observation that can be feed into the output of make_obs_ph
    stochastic: bool
        if set to False all the actions are always deterministic (default False)
    update_eps_ph: float
        update epsilon a new value, if negative not update happens
        (default: no update)

    Returns
    -------
    Tensor of dtype tf.int64 and shape (BATCH_SIZE,) with an action to be performed for
    every element of the batch.


======= train =======

    Function that takes a transition (s,a,r,s') and optimizes Bellman equation's error:

        td_error = Q(s,a) - (r + gamma * max_a' Q(s', a'))
        loss = huber_loss[td_error]

    Parameters
    ----------
    obs_t: object
        a batch of observations
    action: np.array
        actions that were selected upon seeing obs_t.
        dtype must be int32 and shape must be (batch_size,)
    reward: np.array
        immediate reward attained after executing those actions
        dtype must be float32 and shape must be (batch_size,)
    obs_tp1: object
        observations that followed obs_t
    done: np.array
        1 if obs_t was the last observation in the episode and 0 otherwise
        obs_tp1 gets ignored, but must be of the valid shape.
        dtype must be float32 and shape must be (batch_size,)
    weight: np.array
        imporance weights for every element of the batch (gradient is multiplied
        by the importance weight) dtype must be float32 and shape must be (batch_size,)

    Returns
    -------
    td_error: np.array
        a list of differences between Q(s,a) and the target in Bellman's equation.
        dtype is float32 and shape is (batch_size,)

======= update_target ========

    copy the parameters from optimized Q function to the target Q function.
    In Q learning we actually optimize the following error:

        Q(s,a) - (r + gamma * max_a' Q'(s', a'))

    Where Q' is lagging behind Q to stablize the learning. For example for Atari

    Q' is set to Q once every 10000 updates training steps.

"""
import tensorflow as tf
import baselines.common.tf_util as U
import numpy as np


def build_act_mf(make_obs_ph, q_func, z_noise, num_actions, scope="deepq", reuse=None):
    with tf.variable_scope(scope, reuse=reuse):
        observations_ph = U.ensure_tf_input(make_obs_ph("observation"))
        q, q_deterministic, v_mean, v_logvar, z_mean, z_logvar, recon_obs = q_func(observations_ph.get(), z_noise,
                                                                  num_actions,
                                                                  scope="q_func",
                                                                  reuse=tf.AUTO_REUSE)

        act = U.function(inputs=[observations_ph,z_noise],
                         outputs=[z_mean, z_logvar])

        return act


def build_train_mf(make_obs_ph, q_func, num_actions, optimizer, grad_norm_clipping=None, gamma=1.0, scope="mfec",
                   alpha=1.0, beta=1.0, theta=1.0, latent_dim=32, ib=True, reuse=None):
    """Creates the train function:

    Parameters
    ----------
    make_obs_ph: str -> tf.placeholder or TfInput
        a function that takes a name and creates a placeholder of input with that name
    q_func: (tf.Variable, int, str, bool) -> tf.Variable
        the model that takes the following inputs:
            observation_in: object
                the output of observation placeholder
            num_actions: int
                number of actions
            scope: str
            reuse: bool
                should be passed to outer variable scope
        and returns a tensor of shape (batch_size, num_actions) with values of every action.
    num_actions: int
        number of actions
    reuse: bool
        whether or not to reuse the graph variables
    optimizer: tf.train.Optimizer
        optimizer to use for the Q-learning objective.
    grad_norm_clipping: float or None
        clip gradient norms to this value. If None no clipping is performed.
    gamma: float
        discount rate.
    double_q: bool
        if true will use Double Q Learning (https://arxiv.org/abs/1509.06461).
        In general it is a good idea to keep it enabled.
    scope: str or VariableScope
        optional scope for variable_scope.
    reuse: bool or None
        whether or not the variables should be reused. To be able to reuse the scope must be given.

    Returns
    -------
    act: (tf.Variable, bool, float) -> tf.Variable
        function to select and action given observation.
`       See the top of the file for details.
    train: (object, np.array, np.array, object, np.array, np.array) -> np.array
        optimize the error in Bellman's equation.
`       See the top of the file for details.
    update_target: () -> ()
        copy the parameters from optimized Q function to the target Q function.
`       See the top of the file for details.
    debug: {str: function}
        a bunch of functions to print debug data like q_values.
    """
    act_noise = tf.placeholder(tf.float32, [None, latent_dim], name="act_noise")
    act_f = build_act_mf(make_obs_ph, q_func, act_noise, num_actions, scope=scope, reuse=reuse)

    with tf.variable_scope(scope, reuse=reuse):
        # set up placeholders

        # EMDQN

        obs_vae_input = U.ensure_tf_input(make_obs_ph("obs_vae"))
        z_noise_vae = tf.placeholder(tf.float32, [None, latent_dim], name="z_noise_vae")
        inputs = [obs_vae_input,z_noise_vae]
        if ib:
            qec_input = tf.placeholder(tf.float32, [None], name='qec')
            inputs.append(qec_input)
        outputs = []

        q_vae, q_deterministic_vae, v_mean_vae, v_logvar_vae, z_mean_vae, z_logvar_vae, recon_obs = q_func(obs_vae_input.get(),
                                                                                      z_noise_vae, num_actions,
                                                                                      scope="q_func",
                                                                                      reuse=True)
        q_func_vars = U.scope_vars(U.absolute_scope_name("q_func"))

        encoder_loss = -1 + z_mean_vae ** 2 + tf.exp(z_logvar_vae) - z_logvar_vae

        total_loss = tf.reduce_mean(beta * encoder_loss)
        decoder_loss = tf.keras.losses.binary_crossentropy(tf.reshape(recon_obs, [-1]), tf.reshape(
            tf.dtypes.cast(obs_vae_input._placeholder, tf.float32), [-1]))
        print("here", z_mean_vae.shape, z_logvar_vae.shape, encoder_loss.shape, decoder_loss.shape)
        vae_loss = beta * encoder_loss + theta * decoder_loss
        outputs.append(encoder_loss)
        outputs.append(decoder_loss)
        outputs.append(vae_loss)
        total_loss += tf.reduce_mean(theta * decoder_loss)
        if ib:
            ib_loss = (v_mean_vae - tf.stop_gradient(tf.expand_dims(qec_input, 1))) ** 2 / tf.exp(
                v_logvar_vae) + v_logvar_vae
            print("here2", v_mean_vae.shape, tf.expand_dims(qec_input, 1).shape, v_logvar_vae.shape, ib_loss.shape)
            total_ib_loss = alpha * ib_loss + beta * encoder_loss
            outputs.append(total_ib_loss)
            total_loss += tf.reduce_mean(alpha * ib_loss)

        if grad_norm_clipping is not None:
            optimize_expr = U.minimize_and_clip(optimizer,
                                                total_loss,
                                                var_list=q_func_vars,
                                                clip_val=grad_norm_clipping)
        else:
            optimize_expr = optimizer.minimize(total_loss, var_list=q_func_vars)
        # Create callable functions
        # EMDQN
        total_loss_summary = tf.summary.scalar("total loss", total_loss)
        z_var_summary = tf.summary.scalar("z_var", tf.reduce_mean(tf.exp(z_logvar_vae)))
        encoder_loss_summary = tf.summary.scalar("encoder loss", tf.reduce_mean(encoder_loss))
        decoder_loss_summary = tf.summary.scalar("decoder loss", tf.reduce_mean(decoder_loss))
        summaries = [total_loss_summary, z_var_summary, encoder_loss_summary, decoder_loss_summary]
        if ib:
            ib_loss_summary = tf.summary.scalar("ib loss", tf.reduce_mean(ib_loss))
            total_ib_loss_summary = tf.summary.scalar("total ib loss", tf.reduce_mean(total_ib_loss))
            summaries.append(ib_loss_summary)
            summaries.append(total_ib_loss_summary)

        summary = tf.summary.merge(summaries)
        outputs.append(summary)

        train = U.function(
            inputs=inputs,
            outputs=[total_loss,summary],
            updates=[optimize_expr]
        )

        return act_f, train

[((139, 16, 139, 80), 'tensorflow.placeholder', 'tf.placeholder', (), '', True, 'import tensorflow as tf\n'), ((76, 9, 76, 46), 'tensorflow.variable_scope', 'tf.variable_scope', (), '', True, 'import tensorflow as tf\n'), ((83, 14, 84, 52), 'baselines.common.tf_util.function', 'U.function', (), '', True, 'import baselines.common.tf_util as U\n'), ((142, 9, 142, 46), 'tensorflow.variable_scope', 'tf.variable_scope', (), '', True, 'import tensorflow as tf\n'), ((148, 22, 148, 88), 'tensorflow.placeholder', 'tf.placeholder', (), '', True, 'import tensorflow as tf\n'), ((163, 21, 163, 56), 'tensorflow.reduce_mean', 'tf.reduce_mean', ({(163, 36, 163, 55): 'beta * encoder_loss'}, {}), '(beta * encoder_loss)', True, 'import tensorflow as tf\n'), ((171, 22, 171, 58), 'tensorflow.reduce_mean', 'tf.reduce_mean', ({(171, 37, 171, 57): '(theta * decoder_loss)'}, {}), '(theta * decoder_loss)', True, 'import tensorflow as tf\n'), ((189, 29, 189, 72), 'tensorflow.summary.scalar', 'tf.summary.scalar', ({(189, 47, 189, 59): '"""total loss"""', (189, 61, 189, 71): 'total_loss'}, {}), "('total loss', total_loss)", True, 'import tensorflow as tf\n'), ((200, 18, 200, 45), 'tensorflow.summary.merge', 'tf.summary.merge', ({(200, 35, 200, 44): 'summaries'}, {}), '(summaries)', True, 'import tensorflow as tf\n'), ((203, 16, 207, 9), 'baselines.common.tf_util.function', 'U.function', (), '', True, 'import baselines.common.tf_util as U\n'), ((151, 24, 151, 70), 'tensorflow.placeholder', 'tf.placeholder', (), '', True, 'import tensorflow as tf\n'), ((159, 35, 159, 66), 'baselines.common.tf_util.absolute_scope_name', 'U.absolute_scope_name', ({(159, 57, 159, 65): '"""q_func"""'}, {}), "('q_func')", True, 'import baselines.common.tf_util as U\n'), ((164, 59, 164, 86), 'tensorflow.reshape', 'tf.reshape', ({(164, 70, 164, 79): 'recon_obs', (164, 81, 164, 85): '[-1]'}, {}), '(recon_obs, [-1])', True, 'import tensorflow as tf\n'), ((178, 26, 178, 57), 'tensorflow.reduce_mean', 'tf.reduce_mean', ({(178, 41, 178, 56): '(alpha * ib_loss)'}, {}), '(alpha * ib_loss)', True, 'import tensorflow as tf\n'), ((181, 28, 184, 76), 'baselines.common.tf_util.minimize_and_clip', 'U.minimize_and_clip', (), '', True, 'import baselines.common.tf_util as U\n'), ((191, 65, 191, 93), 'tensorflow.reduce_mean', 'tf.reduce_mean', ({(191, 80, 191, 92): 'encoder_loss'}, {}), '(encoder_loss)', True, 'import tensorflow as tf\n'), ((192, 65, 192, 93), 'tensorflow.reduce_mean', 'tf.reduce_mean', ({(192, 80, 192, 92): 'decoder_loss'}, {}), '(decoder_loss)', True, 'import tensorflow as tf\n'), ((161, 46, 161, 66), 'tensorflow.exp', 'tf.exp', ({(161, 53, 161, 65): 'z_logvar_vae'}, {}), '(z_logvar_vae)', True, 'import tensorflow as tf\n'), ((165, 12, 165, 66), 'tensorflow.dtypes.cast', 'tf.dtypes.cast', ({(165, 27, 165, 53): 'obs_vae_input._placeholder', (165, 55, 165, 65): 'tf.float32'}, {}), '(obs_vae_input._placeholder, tf.float32)', True, 'import tensorflow as tf\n'), ((190, 66, 190, 86), 'tensorflow.exp', 'tf.exp', ({(190, 73, 190, 85): 'z_logvar_vae'}, {}), '(z_logvar_vae)', True, 'import tensorflow as tf\n'), ((195, 59, 195, 82), 'tensorflow.reduce_mean', 'tf.reduce_mean', ({(195, 74, 195, 81): 'ib_loss'}, {}), '(ib_loss)', True, 'import tensorflow as tf\n'), ((196, 71, 196, 100), 'tensorflow.reduce_mean', 'tf.reduce_mean', ({(196, 86, 196, 99): 'total_ib_loss'}, {}), '(total_ib_loss)', True, 'import tensorflow as tf\n'), ((173, 91, 174, 29), 'tensorflow.exp', 'tf.exp', ({(174, 16, 174, 28): 'v_logvar_vae'}, {}), '(v_logvar_vae)', True, 'import tensorflow as tf\n'), ((175, 45, 175, 73), 'tensorflow.expand_dims', 'tf.expand_dims', ({(175, 60, 175, 69): 'qec_input', (175, 71, 175, 72): '(1)'}, {}), '(qec_input, 1)', True, 'import tensorflow as tf\n'), ((173, 53, 173, 81), 'tensorflow.expand_dims', 'tf.expand_dims', ({(173, 68, 173, 77): 'qec_input', (173, 79, 173, 80): '(1)'}, {}), '(qec_input, 1)', True, 'import tensorflow as tf\n')]