# Path: disentanglement_lib/data/ground_truth/named_data.py
# def get_named_ground_truth_data(name):
# 
# Path: disentanglement_lib/postprocessing/methods.py
# def mean_representation(
#     ground_truth_data,
#     gaussian_encoder,
#     random_state,
#     save_path,
# ):
#   def transform_fn(mean, logvar):
# def sampled_representation(ground_truth_data, gaussian_encoder, random_state,
#                            save_path):
#   def transform_fn(mean, logvar):
# 
# Path: disentanglement_lib/utils/convolute_hub.py
# def convolute_and_save(module_path, signature, export_path, transform_fn,
#                        transform_checkpoint_path, new_signature=None):
#   def module_fn():
# def save_numpy_arrays_to_checkpoint(checkpoint_path, **dict_with_arrays):
# def _placeholders_from_module(tfhub_module, signature):
# 
# 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):


import os
import time
import numpy as np
import tensorflow.compat.v1 as tf
import tensorflow_hub as hub
import gin.tf
from disentanglement_lib.data.ground_truth import named_data
from disentanglement_lib.postprocessing import methods  # pylint: disable=unused-import
from disentanglement_lib.utils import convolute_hub
from disentanglement_lib.utils import results
  # Path to TFHub module of previously trained model.
  module_path = os.path.join(model_dir, "tfhub")
  with hub.eval_function_for_module(module_path) as f:

    def _gaussian_encoder(x):
      """Encodes images using trained model."""
      # Push images through the TFHub module.
      output = f(dict(images=x), signature="gaussian_encoder", as_dict=True)
      # Convert to numpy arrays and return.
      return {key: np.array(values) for key, values in output.items()}

    # Run the postprocessing function which returns a transformation function
    # that can be used to create the representation from the mean and log
    # variance of the Gaussian distribution given by the encoder. Also returns
    # path to a checkpoint if the transformation requires variables.
    transform_fn, transform_checkpoint_path = postprocess_fn(
        dataset, _gaussian_encoder, np.random.RandomState(random_seed),
        output_dir)

    # Takes the "gaussian_encoder" signature, extracts the representation and
    # then saves under the signature "representation".
    tfhub_module_dir = os.path.join(output_dir, "tfhub")
    convolute_hub.convolute_and_save(
        module_path, "gaussian_encoder", tfhub_module_dir, transform_fn,
        transform_checkpoint_path, "representation")

  # We first copy over all the prior results and configs.
  original_results_dir = os.path.join(model_dir, "results")
  results_dir = os.path.join(output_dir, "results")
  results_dict = dict(elapsed_time=time.time() - experiment_timer)

# Path: disentanglement_lib/data/ground_truth/named_data.py
# def get_named_ground_truth_data(name):
# 
# Path: disentanglement_lib/postprocessing/methods.py
# def mean_representation(
#     ground_truth_data,
#     gaussian_encoder,
#     random_state,
#     save_path,
# ):
#   def transform_fn(mean, logvar):
# def sampled_representation(ground_truth_data, gaussian_encoder, random_state,
#                            save_path):
#   def transform_fn(mean, logvar):
# 
# Path: disentanglement_lib/utils/convolute_hub.py
# def convolute_and_save(module_path, signature, export_path, transform_fn,
#                        transform_checkpoint_path, new_signature=None):
#   def module_fn():
# def save_numpy_arrays_to_checkpoint(checkpoint_path, **dict_with_arrays):
# def _placeholders_from_module(tfhub_module, signature):
# 
# 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/postprocessing/postprocess.py
import os
import time
import numpy as np
import tensorflow.compat.v1 as tf
import tensorflow_hub as hub
import gin.tf
from disentanglement_lib.data.ground_truth import named_data
from disentanglement_lib.postprocessing import methods  # pylint: disable=unused-import
from disentanglement_lib.utils import convolute_hub
from disentanglement_lib.utils import results

  # Path to TFHub module of previously trained model.
  module_path = os.path.join(model_dir, "tfhub")
  with hub.eval_function_for_module(module_path) as f:

    def _gaussian_encoder(x):
      """Encodes images using trained model."""
      # Push images through the TFHub module.
      output = f(dict(images=x), signature="gaussian_encoder", as_dict=True)
      # Convert to numpy arrays and return.
      return {key: np.array(values) for key, values in output.items()}

    # Run the postprocessing function which returns a transformation function
    # that can be used to create the representation from the mean and log
    # variance of the Gaussian distribution given by the encoder. Also returns
    # path to a checkpoint if the transformation requires variables.
    transform_fn, transform_checkpoint_path = postprocess_fn(
        dataset, _gaussian_encoder, np.random.RandomState(random_seed),
        output_dir)

    # Takes the "gaussian_encoder" signature, extracts the representation and
    # then saves under the signature "representation".
    tfhub_module_dir = os.path.join(output_dir, "tfhub")
    convolute_hub.convolute_and_save(
        module_path, "gaussian_encoder", tfhub_module_dir, transform_fn,
        transform_checkpoint_path, "representation")

  # We first copy over all the prior results and configs.
  original_results_dir = os.path.join(model_dir, "results")
  results_dir = os.path.join(output_dir, "results")
  results_dict = dict(elapsed_time=time.time() - experiment_timer)

  results.update_result_directory(results_dir, "postprocess", results_dict,