# Path: disentanglement_lib/methods/unsupervised/vae.py
# class BaseVAE(gaussian_encoder_model.GaussianEncoderModel):
# class BetaVAE(BaseVAE):
# class AnnealedVAE(BaseVAE):
# class FactorVAE(BaseVAE):
# class DIPVAE(BaseVAE):
# class BetaTCVAE(BaseVAE):
#   def model_fn(self, features, labels, mode, params):
#   def gaussian_encoder(self, input_tensor, is_training):
#   def decode(self, latent_tensor, observation_shape, is_training):
# def shuffle_codes(z):
# def compute_gaussian_kl(z_mean, z_logvar):
# def make_metric_fn(*names):
#   def metric_fn(*args):
#   def __init__(self, beta=gin.REQUIRED):
#   def regularizer(self, kl_loss, z_mean, z_logvar, z_sampled):
# def anneal(c_max, step, iteration_threshold):
#   def __init__(self,
#                gamma=gin.REQUIRED,
#                c_max=gin.REQUIRED,
#                iteration_threshold=gin.REQUIRED):
#   def regularizer(self, kl_loss, z_mean, z_logvar, z_sampled):
#   def __init__(self, gamma=gin.REQUIRED):
#   def model_fn(self, features, labels, mode, params):
# def compute_covariance_z_mean(z_mean):
# def regularize_diag_off_diag_dip(covariance_matrix, lambda_od, lambda_d):
#   def __init__(self,
#                lambda_od=gin.REQUIRED,
#                lambda_d_factor=gin.REQUIRED,
#                dip_type="i"):
#   def regularizer(self, kl_loss, z_mean, z_logvar, z_sampled):
# def gaussian_log_density(samples, mean, log_var):
# def total_correlation(z, z_mean, z_logvar):
#   def __init__(self, beta=gin.REQUIRED):
#   def regularizer(self, kl_loss, z_mean, z_logvar, z_sampled):


from absl.testing import parameterized
from disentanglement_lib.methods.unsupervised import vae
import numpy as np
import tensorflow.compat.v1 as tf
import gin.tf.external_configurables  # pylint: disable=unused-import

  @parameterized.parameters(
      (True, 0., 1.), (True, 0., 4.), (True, 1., 1.),
      (False, np.zeros(10), np.ones([10, 10])),
      (False, np.zeros(10), _make_symmetric_psd(np.random.random((10, 10)))))
  def test_compute_covariance_z_mean(self, diagonal, mean, cov):
    if diagonal:
      samples = tf.random.normal(
          shape=(100000, 10), mean=mean, stddev=tf.math.sqrt(cov))
      cov = np.diag(np.ones([10])) * cov
    else:
      samples = tf.constant(
          np.random.multivariate_normal(mean, cov, size=(1000000)))
    with self.test_session() as sess:
      test_value = sess.run(vae.compute_covariance_z_mean(samples))
      self.assertBetween(np.sum((test_value - cov)**2), 0., 0.1)

  @parameterized.parameters(
      (np.ones([10, 10]), 90., 90.1), (np.zeros([10, 10]), 10., 10.1),
      (np.diag(np.ones(10)), 0., 0.1), (2. * np.diag(np.ones(10)), 10., 10.1))
  def test_regularize_diag_off_diag_dip(self, matrix, target_low, target_high):
    matrix_tf = tf.convert_to_tensor(matrix, dtype=np.float32)
    with self.test_session() as sess:
      test_value = sess.run(vae.regularize_diag_off_diag_dip(matrix_tf, 1, 1))
      self.assertBetween(test_value, target_low, target_high)

  @parameterized.parameters((0., -1.4190, -1.4188), (1., -0.92, -0.91))
  def test_gaussian_log_density(self, z_mean, target_low, target_high):
    matrix = tf.ones(1)
    with self.test_session() as sess:

# Path: disentanglement_lib/methods/unsupervised/vae.py
# class BaseVAE(gaussian_encoder_model.GaussianEncoderModel):
# class BetaVAE(BaseVAE):
# class AnnealedVAE(BaseVAE):
# class FactorVAE(BaseVAE):
# class DIPVAE(BaseVAE):
# class BetaTCVAE(BaseVAE):
#   def model_fn(self, features, labels, mode, params):
#   def gaussian_encoder(self, input_tensor, is_training):
#   def decode(self, latent_tensor, observation_shape, is_training):
# def shuffle_codes(z):
# def compute_gaussian_kl(z_mean, z_logvar):
# def make_metric_fn(*names):
#   def metric_fn(*args):
#   def __init__(self, beta=gin.REQUIRED):
#   def regularizer(self, kl_loss, z_mean, z_logvar, z_sampled):
# def anneal(c_max, step, iteration_threshold):
#   def __init__(self,
#                gamma=gin.REQUIRED,
#                c_max=gin.REQUIRED,
#                iteration_threshold=gin.REQUIRED):
#   def regularizer(self, kl_loss, z_mean, z_logvar, z_sampled):
#   def __init__(self, gamma=gin.REQUIRED):
#   def model_fn(self, features, labels, mode, params):
# def compute_covariance_z_mean(z_mean):
# def regularize_diag_off_diag_dip(covariance_matrix, lambda_od, lambda_d):
#   def __init__(self,
#                lambda_od=gin.REQUIRED,
#                lambda_d_factor=gin.REQUIRED,
#                dip_type="i"):
#   def regularizer(self, kl_loss, z_mean, z_logvar, z_sampled):
# def gaussian_log_density(samples, mean, log_var):
# def total_correlation(z, z_mean, z_logvar):
#   def __init__(self, beta=gin.REQUIRED):
#   def regularizer(self, kl_loss, z_mean, z_logvar, z_sampled):


# Path: disentanglement_lib/methods/unsupervised/vae_test.py
from absl.testing import parameterized
from disentanglement_lib.methods.unsupervised import vae
import numpy as np
import tensorflow.compat.v1 as tf
import gin.tf.external_configurables  # pylint: disable=unused-import


  @parameterized.parameters(
      (True, 0., 1.), (True, 0., 4.), (True, 1., 1.),
      (False, np.zeros(10), np.ones([10, 10])),
      (False, np.zeros(10), _make_symmetric_psd(np.random.random((10, 10)))))
  def test_compute_covariance_z_mean(self, diagonal, mean, cov):
    if diagonal:
      samples = tf.random.normal(
          shape=(100000, 10), mean=mean, stddev=tf.math.sqrt(cov))
      cov = np.diag(np.ones([10])) * cov
    else:
      samples = tf.constant(
          np.random.multivariate_normal(mean, cov, size=(1000000)))
    with self.test_session() as sess:
      test_value = sess.run(vae.compute_covariance_z_mean(samples))
      self.assertBetween(np.sum((test_value - cov)**2), 0., 0.1)

  @parameterized.parameters(
      (np.ones([10, 10]), 90., 90.1), (np.zeros([10, 10]), 10., 10.1),
      (np.diag(np.ones(10)), 0., 0.1), (2. * np.diag(np.ones(10)), 10., 10.1))
  def test_regularize_diag_off_diag_dip(self, matrix, target_low, target_high):
    matrix_tf = tf.convert_to_tensor(matrix, dtype=np.float32)
    with self.test_session() as sess:
      test_value = sess.run(vae.regularize_diag_off_diag_dip(matrix_tf, 1, 1))
      self.assertBetween(test_value, target_low, target_high)

  @parameterized.parameters((0., -1.4190, -1.4188), (1., -0.92, -0.91))
  def test_gaussian_log_density(self, z_mean, target_low, target_high):
    matrix = tf.ones(1)
    with self.test_session() as sess:

      test_value = sess.run(vae.gaussian_log_density(matrix, z_mean, 0.))[0]