# Path: disentanglement_lib/data/ground_truth/named_data.py
# def get_named_ground_truth_data(name):
# 
# Path: disentanglement_lib/methods/unsupervised/gaussian_encoder_model.py
# class GaussianEncoderModel(object):
#   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 sample_from_latent_distribution(self, z_mean, z_logvar):
# def export_as_tf_hub(gaussian_encoder_model,
#                      observation_shape,
#                      checkpoint_path,
#                      export_path,
#                      drop_collections=None):
#   def module_fn(is_training):
# 
# Path: disentanglement_lib/methods/weak/weak_vae.py
# def make_weak_loss(z1, z2, labels, loss_fn=gin.REQUIRED):
#   def __init__(self, beta=gin.REQUIRED):
#   def regularizer(self, kl_loss, z_mean, z_logvar, z_sampled):
#   def model_fn(self, features, labels, mode, params):
#   def aggregate(self, z_mean, z_logvar, new_mean, new_log_var, labels,
#                 kl_per_point):
#   def aggregate(self, z_mean, z_logvar, new_mean, new_log_var, labels,
#                 kl_per_point):
#   def __init__(self, beta=gin.REQUIRED):
#   def regularizer(self, kl_loss, z_mean, z_logvar, z_sampled):
#   def model_fn(self, features, labels, mode, params):
#   def aggregate(self, z_mean, z_logvar, new_mean, new_log_var, labels,
#                 kl_per_point):
#   def aggregate(self, z_mean, z_logvar, new_mean, new_log_var, labels,
#                 kl_per_point):
# def aggregate_labels(z_mean, z_logvar, new_mean, new_log_var, labels,
#                      kl_per_point):
# def aggregate_argmax(z_mean, z_logvar, new_mean, new_log_var, labels,
#                      kl_per_point):
# def discretize_in_bins(x):
# def compute_kl(z_1, z_2, logvar_1, logvar_2):
# def make_metric_fn(*names):
#   def metric_fn(*args):
# class GroupVAEBase(vae.BaseVAE):
# class GroupVAELabels(GroupVAEBase):
# class GroupVAEArgmax(GroupVAEBase):
# class MLVae(vae.BaseVAE):
# class MLVaeLabels(MLVae):
# class MLVaeArgmax(MLVae):
# 
# Path: disentanglement_lib/utils/results.py
# def update_result_directory(result_directory,
#                             step_name,
#                             results_dict,
#                             old_result_directory=None):
# def _copy_recursively(path_to_old_dir, path_to_new_dir):
# def copydir(path_to_old_dir, path_to_new_dir):
# def save_gin(config_path):
#   def default(self, obj):
# def save_dict(config_path, dict_with_info):
# def gin_dict(config_path=None):
# def namespaced_dict(base_dict=None, **named_dicts):
# def aggregate_json_results(base_path):
# class Encoder(json.JSONEncoder):
# 
# Path: disentanglement_lib/visualize/visualize_util.py
# def save_image(image, image_path):
# def grid_save_images(images, image_path):
# def padded_grid(images, num_rows=None, padding_px=10, value=None):
# def padded_stack(images, padding_px=10, axis=0, value=None):
# def padding_array(image, padding_px, axis, value=None):
# def best_num_rows(num_elements, max_ratio=4):
# def pad_around(image, padding_px=10, axis=None, value=None):
# def add_below(image, padding_px=10, value=None):
# def save_animation(list_of_animated_images, image_path, fps):
# def cycle_factor(starting_index, num_indices, num_frames):
# def cycle_gaussian(starting_value, num_frames, loc=0., scale=1.):
# def cycle_interval(starting_value, num_frames, min_val, max_val):


import os
import time
import numpy as np
import tensorflow as tf
import gin.tf.external_configurables  # pylint: disable=unused-import
import gin.tf
from disentanglement_lib.data.ground_truth import named_data
from disentanglement_lib.methods.unsupervised import gaussian_encoder_model
from disentanglement_lib.methods.weak import weak_vae  # pylint: disable=unused-import
from disentanglement_lib.utils import results
from disentanglement_lib.visualize import visualize_util
from tensorflow_estimator.python.estimator.tpu import tpu_config
from tensorflow_estimator.python.estimator.tpu.tpu_estimator import TPUEstimator
  """Trains the estimator and exports the snapshot and the gin config.

  The use of this function requires the gin binding 'dataset.name' to be
  specified as that determines the data set used for training.

  Args:
    model_dir: String with path to directory where model output should be saved.
    overwrite: Boolean indicating whether to overwrite output directory.
    model: GaussianEncoderModel that should be trained and exported.
    training_steps: Integer with number of training steps.
    random_seed: Integer with random seed used for training.
    batch_size: Integer with the batch size.
    name: Optional string with name of the model (can be used to name models).
  """
  # We do not use the variable 'name'. Instead, it can be used to name results
  # as it will be part of the saved gin config.
  del name

  # Delete the output directory if necessary.
  if tf.compat.v1.gfile.IsDirectory(model_dir):
    if overwrite:
      tf.compat.v1.gfile.DeleteRecursively(model_dir)
    else:
      raise ValueError("Directory already exists and overwrite is False.")

  # Create a numpy random state. We will sample the random seeds for training
  # and evaluation from this.
  random_state = np.random.RandomState(random_seed)

  # Obtain the dataset.

# Path: disentanglement_lib/data/ground_truth/named_data.py
# def get_named_ground_truth_data(name):
# 
# Path: disentanglement_lib/methods/unsupervised/gaussian_encoder_model.py
# class GaussianEncoderModel(object):
#   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 sample_from_latent_distribution(self, z_mean, z_logvar):
# def export_as_tf_hub(gaussian_encoder_model,
#                      observation_shape,
#                      checkpoint_path,
#                      export_path,
#                      drop_collections=None):
#   def module_fn(is_training):
# 
# Path: disentanglement_lib/methods/weak/weak_vae.py
# def make_weak_loss(z1, z2, labels, loss_fn=gin.REQUIRED):
#   def __init__(self, beta=gin.REQUIRED):
#   def regularizer(self, kl_loss, z_mean, z_logvar, z_sampled):
#   def model_fn(self, features, labels, mode, params):
#   def aggregate(self, z_mean, z_logvar, new_mean, new_log_var, labels,
#                 kl_per_point):
#   def aggregate(self, z_mean, z_logvar, new_mean, new_log_var, labels,
#                 kl_per_point):
#   def __init__(self, beta=gin.REQUIRED):
#   def regularizer(self, kl_loss, z_mean, z_logvar, z_sampled):
#   def model_fn(self, features, labels, mode, params):
#   def aggregate(self, z_mean, z_logvar, new_mean, new_log_var, labels,
#                 kl_per_point):
#   def aggregate(self, z_mean, z_logvar, new_mean, new_log_var, labels,
#                 kl_per_point):
# def aggregate_labels(z_mean, z_logvar, new_mean, new_log_var, labels,
#                      kl_per_point):
# def aggregate_argmax(z_mean, z_logvar, new_mean, new_log_var, labels,
#                      kl_per_point):
# def discretize_in_bins(x):
# def compute_kl(z_1, z_2, logvar_1, logvar_2):
# def make_metric_fn(*names):
#   def metric_fn(*args):
# class GroupVAEBase(vae.BaseVAE):
# class GroupVAELabels(GroupVAEBase):
# class GroupVAEArgmax(GroupVAEBase):
# class MLVae(vae.BaseVAE):
# class MLVaeLabels(MLVae):
# class MLVaeArgmax(MLVae):
# 
# Path: disentanglement_lib/utils/results.py
# def update_result_directory(result_directory,
#                             step_name,
#                             results_dict,
#                             old_result_directory=None):
# def _copy_recursively(path_to_old_dir, path_to_new_dir):
# def copydir(path_to_old_dir, path_to_new_dir):
# def save_gin(config_path):
#   def default(self, obj):
# def save_dict(config_path, dict_with_info):
# def gin_dict(config_path=None):
# def namespaced_dict(base_dict=None, **named_dicts):
# def aggregate_json_results(base_path):
# class Encoder(json.JSONEncoder):
# 
# Path: disentanglement_lib/visualize/visualize_util.py
# def save_image(image, image_path):
# def grid_save_images(images, image_path):
# def padded_grid(images, num_rows=None, padding_px=10, value=None):
# def padded_stack(images, padding_px=10, axis=0, value=None):
# def padding_array(image, padding_px, axis, value=None):
# def best_num_rows(num_elements, max_ratio=4):
# def pad_around(image, padding_px=10, axis=None, value=None):
# def add_below(image, padding_px=10, value=None):
# def save_animation(list_of_animated_images, image_path, fps):
# def cycle_factor(starting_index, num_indices, num_frames):
# def cycle_gaussian(starting_value, num_frames, loc=0., scale=1.):
# def cycle_interval(starting_value, num_frames, min_val, max_val):


# Path: disentanglement_lib/methods/weak/train_weak_lib.py
import os
import time
import numpy as np
import tensorflow as tf
import gin.tf.external_configurables  # pylint: disable=unused-import
import gin.tf
from disentanglement_lib.data.ground_truth import named_data
from disentanglement_lib.methods.unsupervised import gaussian_encoder_model
from disentanglement_lib.methods.weak import weak_vae  # pylint: disable=unused-import
from disentanglement_lib.utils import results
from disentanglement_lib.visualize import visualize_util
from tensorflow_estimator.python.estimator.tpu import tpu_config
from tensorflow_estimator.python.estimator.tpu.tpu_estimator import TPUEstimator

  """Trains the estimator and exports the snapshot and the gin config.

  The use of this function requires the gin binding 'dataset.name' to be
  specified as that determines the data set used for training.

  Args:
    model_dir: String with path to directory where model output should be saved.
    overwrite: Boolean indicating whether to overwrite output directory.
    model: GaussianEncoderModel that should be trained and exported.
    training_steps: Integer with number of training steps.
    random_seed: Integer with random seed used for training.
    batch_size: Integer with the batch size.
    name: Optional string with name of the model (can be used to name models).
  """
  # We do not use the variable 'name'. Instead, it can be used to name results
  # as it will be part of the saved gin config.
  del name

  # Delete the output directory if necessary.
  if tf.compat.v1.gfile.IsDirectory(model_dir):
    if overwrite:
      tf.compat.v1.gfile.DeleteRecursively(model_dir)
    else:
      raise ValueError("Directory already exists and overwrite is False.")

  # Create a numpy random state. We will sample the random seeds for training
  # and evaluation from this.
  random_state = np.random.RandomState(random_seed)

  # Obtain the dataset.

  dataset = named_data.get_named_ground_truth_data()