import tensorflow as tf

import numpy as np
import tensorflow as tf

class VariationalDense:
    """Variational Dense Layer Class"""
    def __init__(self, n_in, n_out, dropout_mask_ph,
                 model_prob=0.9, model_lam=3e-4, activation=None, name="hidden"):
        self.model_prob = model_prob    # probability to keep units
        self.model_lam = model_lam      # l^2 / 2*tau: l=1e-2, tau=[0.1, 0.15, 0.2]

        self.dropout_mask_ph = dropout_mask_ph # placeholder: p_s * i_s
        self.p_s = tf.shape(self.dropout_mask_ph)[0] # post sample size
        self.DM = tf.zeros(shape=[self.p_s, n_in, n_in]) # Dropout masks: p_s * i_s * i_s
        self.DM = tf.linalg.set_diag(self.DM, self.dropout_mask_ph)

        kernel_initializer = tf.initializers.truncated_normal(mean=0.0, stddev=0.01)
        self.model_W = tf.get_variable("{}_W".format(name), initializer=kernel_initializer([n_in, n_out])) # variational parameters
        self.model_b = tf.get_variable("{}_b".format(name), initializer=tf.zeros([n_out]))

        self.model_DMW = tf.einsum('pij,jk->pik', self.DM, self.model_W) # Masked weight: p_s * i_s * o_s
        self.model_tiled_b = tf.tile(tf.reshape(self.model_b, [1, n_out]), [self.p_s, 1])

        if activation is None:
            self.activation = tf.identity
        else:
            self.activation = activation
