import tensorflow as tf


# TODO: add batch normalization
def mlp_variational(x, dropout_mask_phs, hidden_sizes=(32,),
                    activation=tf.tanh, output_activation=None, dropout_rate=0.1):

    # layer_sizes = (input_size, h1, h2, ..., output_size)
    layer_sizes = hidden_sizes.copy()
    layer_sizes.insert(0, x.shape.as_list()[1])

    # tile x from shape (b_s * i_s) to (p_s * b_s * i_s)
    post_size = tf.shape(dropout_mask_phs[0])[0]
    x = tf.tile(tf.reshape(x, [1, tf.shape(x)[0], tf.shape(x)[1]]), [post_size, 1, 1])
    # TODO: no dropout on input
    regularization = 0
    # Create hidden layers
    for layer_i in range(1,len(layer_sizes)-1):
        hidden_layer = VariationalDense(n_in=layer_sizes[layer_i-1],
                                        n_out=layer_sizes[layer_i],
                                        dropout_mask_ph=dropout_mask_phs[layer_i-1],
                                        model_prob=1.0 - dropout_rate,
