# Path: disentanglement_lib/evaluation/abstract_reasoning/models.py
# class TwoStageModel(object):
# class BaselineCNNEmbedder(tf.keras.Model):
# class HubEmbedding(tf.keras.Model):
# class OptimizedWildRelNet(tf.keras.Model):
#   def __init__(self,
#                embedding_model_class=gin.REQUIRED,
#                reasoning_model_class=gin.REQUIRED,
#                optimizer_fn=None):
#   def model_fn(self, features, labels, mode, params):
#       def metric_fn(labels, logits):
#   def __init__(self,
#                num_latent=gin.REQUIRED,
#                name="BaselineCNNEmbedder",
#                **kwargs):
#   def call(self, inputs, **kwargs):
#   def __init__(self, hub_path=gin.REQUIRED, name="HubEmbedding", **kwargs):
#     def _embedder(x):
#   def call(self, inputs, **kwargs):
#   def __init__(self,
#                edge_mlp=gin.REQUIRED,
#                graph_mlp=gin.REQUIRED,
#                dropout_in_last_graph_layer=gin.REQUIRED,
#                name="OptimizedWildRelNet",
#                **kwargs):
#   def call(self, inputs, **kwargs):
# def get_activation(activation=tf.keras.activations.relu):
# def get_kernel_initializer(kernel_initializer="lecun_normal"):
# 
# Path: disentanglement_lib/evaluation/abstract_reasoning/pgm_data.py
# def get_pgm_dataset(pgm_type=gin.REQUIRED):
#   def __init__(self, ground_truth_data, sampling_strategy, relations_dist):
#   def sample(self, random_state):
#   def tf_data_set(self, seed):
#     def generator():
#   def make_input_fn(self, seed, num_batches=None):
#     def input_fn(params):
#   def __init__(self,
#                solution,
#                alternatives,
#                position,
#                solution_factors=None,
#                alternatives_factors=None,
#                num_factor_values=None):
#   def get_context(self):
#   def get_answers(self):
#   def get_context_factor_values(self):
#   def get_answers_factor_values(self):
#   def range_embed_factors(self, factors):
#   def onehot_embed_factors(self, factors):
#   def training_sample(self):
#   def make_image(self, answer=False, padding_px=8, border_px=4):
#   def __init__(self, wrapped_ground_truth_data, max_factors):
#   def num_factors(self):
#   def factors_num_values(self):
#   def observation_shape(self):
#   def sample_factors(self, num, random_state):
#   def _sample_factor(self, i, num, random_state):
#   def sample_observations_from_factors(self, factors, random_state):
# def question_mark():
# def onehot(indices, num_atoms):
# class PGMDataset(object):
# class PGMInstance(object):
# class Quantizer(gtd.GroundTruthData):
# COLORS = {
#     "blue": np.array([66., 103., 210.]) / 255.,
#     "red": np.array([234., 67., 53.]) / 255.,
#     "yellow": np.array([251., 188., 4.]) / 255.,
#     "green": np.array([52., 168., 83.]) / 255.,
#     "grey": np.array([154., 160., 166.]) / 255.,
# }
# QUESTION_MARK = [None]
#       QUESTION_MARK[0] = np.array(Image.open(f).convert("RGB")) * 1.0 / 255.
# 
# 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 gin.tf.external_configurables  # pylint: disable=unused-import
import gin.tf
from disentanglement_lib.evaluation.abstract_reasoning import models  # pylint: disable=unused-import
from disentanglement_lib.evaluation.abstract_reasoning import pgm_data
from disentanglement_lib.utils import results
from tensorflow.contrib import tpu as contrib_tpu
        input_fn=dataset.make_input_fn(random_state.randint(2**32)),
        steps=training_steps_per_iteration)
    # Compute validation scores used for model selection.
    validation_results = tpu_estimator.evaluate(
        input_fn=dataset.make_input_fn(
            random_state.randint(2**32), num_batches=eval_steps_per_iteration))
    validation_scores.append(validation_results["accuracy"])
    tf.logging.info("Validation results %s", validation_results)
    # Compute test scores for final results.
    test_results = tpu_estimator.evaluate(
        input_fn=dataset.make_input_fn(
            random_state.randint(2**32), num_batches=eval_steps_per_iteration),
        name="test")
    dict_at_iteration = results.namespaced_dict(
        val=validation_results, test=test_results)
    results_dict_of_dicts["step{}".format(steps_so_far)] = dict_at_iteration
    all_dicts.append(dict_at_iteration)

  # Select the best number of steps based on the validation scores and add it as
  # as a special key to the dictionary.
  best_index = np.argmax(validation_scores)
  results_dict_of_dicts["best"] = all_dicts[best_index]

  # Save the results. The result dir will contain all the results and config
  # files that we copied along, as we progress in the pipeline. The idea is that
  # these files will be available for analysis at the end.
  if input_dir is not None:
    original_results_dir = os.path.join(input_dir, "results")
  else:
    original_results_dir = None

# Path: disentanglement_lib/evaluation/abstract_reasoning/models.py
# class TwoStageModel(object):
# class BaselineCNNEmbedder(tf.keras.Model):
# class HubEmbedding(tf.keras.Model):
# class OptimizedWildRelNet(tf.keras.Model):
#   def __init__(self,
#                embedding_model_class=gin.REQUIRED,
#                reasoning_model_class=gin.REQUIRED,
#                optimizer_fn=None):
#   def model_fn(self, features, labels, mode, params):
#       def metric_fn(labels, logits):
#   def __init__(self,
#                num_latent=gin.REQUIRED,
#                name="BaselineCNNEmbedder",
#                **kwargs):
#   def call(self, inputs, **kwargs):
#   def __init__(self, hub_path=gin.REQUIRED, name="HubEmbedding", **kwargs):
#     def _embedder(x):
#   def call(self, inputs, **kwargs):
#   def __init__(self,
#                edge_mlp=gin.REQUIRED,
#                graph_mlp=gin.REQUIRED,
#                dropout_in_last_graph_layer=gin.REQUIRED,
#                name="OptimizedWildRelNet",
#                **kwargs):
#   def call(self, inputs, **kwargs):
# def get_activation(activation=tf.keras.activations.relu):
# def get_kernel_initializer(kernel_initializer="lecun_normal"):
# 
# Path: disentanglement_lib/evaluation/abstract_reasoning/pgm_data.py
# def get_pgm_dataset(pgm_type=gin.REQUIRED):
#   def __init__(self, ground_truth_data, sampling_strategy, relations_dist):
#   def sample(self, random_state):
#   def tf_data_set(self, seed):
#     def generator():
#   def make_input_fn(self, seed, num_batches=None):
#     def input_fn(params):
#   def __init__(self,
#                solution,
#                alternatives,
#                position,
#                solution_factors=None,
#                alternatives_factors=None,
#                num_factor_values=None):
#   def get_context(self):
#   def get_answers(self):
#   def get_context_factor_values(self):
#   def get_answers_factor_values(self):
#   def range_embed_factors(self, factors):
#   def onehot_embed_factors(self, factors):
#   def training_sample(self):
#   def make_image(self, answer=False, padding_px=8, border_px=4):
#   def __init__(self, wrapped_ground_truth_data, max_factors):
#   def num_factors(self):
#   def factors_num_values(self):
#   def observation_shape(self):
#   def sample_factors(self, num, random_state):
#   def _sample_factor(self, i, num, random_state):
#   def sample_observations_from_factors(self, factors, random_state):
# def question_mark():
# def onehot(indices, num_atoms):
# class PGMDataset(object):
# class PGMInstance(object):
# class Quantizer(gtd.GroundTruthData):
# COLORS = {
#     "blue": np.array([66., 103., 210.]) / 255.,
#     "red": np.array([234., 67., 53.]) / 255.,
#     "yellow": np.array([251., 188., 4.]) / 255.,
#     "green": np.array([52., 168., 83.]) / 255.,
#     "grey": np.array([154., 160., 166.]) / 255.,
# }
# QUESTION_MARK = [None]
#       QUESTION_MARK[0] = np.array(Image.open(f).convert("RGB")) * 1.0 / 255.
# 
# 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/evaluation/abstract_reasoning/reason.py
import os
import time
import numpy as np
import tensorflow.compat.v1 as tf
import gin.tf.external_configurables  # pylint: disable=unused-import
import gin.tf
from disentanglement_lib.evaluation.abstract_reasoning import models  # pylint: disable=unused-import
from disentanglement_lib.evaluation.abstract_reasoning import pgm_data
from disentanglement_lib.utils import results
from tensorflow.contrib import tpu as contrib_tpu

        input_fn=dataset.make_input_fn(random_state.randint(2**32)),
        steps=training_steps_per_iteration)
    # Compute validation scores used for model selection.
    validation_results = tpu_estimator.evaluate(
        input_fn=dataset.make_input_fn(
            random_state.randint(2**32), num_batches=eval_steps_per_iteration))
    validation_scores.append(validation_results["accuracy"])
    tf.logging.info("Validation results %s", validation_results)
    # Compute test scores for final results.
    test_results = tpu_estimator.evaluate(
        input_fn=dataset.make_input_fn(
            random_state.randint(2**32), num_batches=eval_steps_per_iteration),
        name="test")
    dict_at_iteration = results.namespaced_dict(
        val=validation_results, test=test_results)
    results_dict_of_dicts["step{}".format(steps_so_far)] = dict_at_iteration
    all_dicts.append(dict_at_iteration)

  # Select the best number of steps based on the validation scores and add it as
  # as a special key to the dictionary.
  best_index = np.argmax(validation_scores)
  results_dict_of_dicts["best"] = all_dicts[best_index]

  # Save the results. The result dir will contain all the results and config
  # files that we copied along, as we progress in the pipeline. The idea is that
  # these files will be available for analysis at the end.
  if input_dir is not None:
    original_results_dir = os.path.join(input_dir, "results")
  else:
    original_results_dir = None

  results_dict = results.namespaced_dict(**results_dict_of_dicts)