# Path: disentanglement_lib/evaluation/metrics/dci.py
# def compute_dci(ground_truth_data, representation_function, random_state,
#                 artifact_dir=None,
#                 num_train=gin.REQUIRED,
#                 num_test=gin.REQUIRED,
#                 batch_size=16):
# def _compute_dci(mus_train, ys_train, mus_test, ys_test):
# def compute_dci_on_fixed_data(observations, labels, representation_function,
#                               train_percentage=gin.REQUIRED, batch_size=100):
# def compute_importance_gbt(x_train, y_train, x_test, y_test):
# def disentanglement_per_code(importance_matrix):
# def disentanglement(importance_matrix):
# def completeness_per_factor(importance_matrix):
# def completeness(importance_matrix):
# 
# Path: disentanglement_lib/evaluation/metrics/modularity_explicitness.py
# def compute_modularity_explicitness(ground_truth_data,
#                                     representation_function,
#                                     random_state,
#                                     artifact_dir=None,
#                                     num_train=gin.REQUIRED,
#                                     num_test=gin.REQUIRED,
#                                     batch_size=16):
# def explicitness_per_factor(mus_train, y_train, mus_test, y_test):
# def modularity(mutual_information):
# 
# Path: disentanglement_lib/evaluation/metrics/sap_score.py
# def compute_sap(ground_truth_data,
#                 representation_function,
#                 random_state,
#                 artifact_dir=None,
#                 num_train=gin.REQUIRED,
#                 num_test=gin.REQUIRED,
#                 batch_size=16,
#                 continuous_factors=gin.REQUIRED):
# def _compute_sap(mus, ys, mus_test, ys_test, continuous_factors):
# def compute_sap_on_fixed_data(observations, labels, representation_function,
#                               train_percentage=gin.REQUIRED,
#                               continuous_factors=gin.REQUIRED,
#                               batch_size=100):
# def compute_score_matrix(mus, ys, mus_test, ys_test, continuous_factors):
# def compute_avg_diff_top_two(matrix):
# 
# Path: disentanglement_lib/evaluation/metrics/utils.py
# def generate_batch_factor_code(ground_truth_data, representation_function,
#                                num_points, random_state, batch_size):
# def split_train_test(observations, train_percentage):
# def obtain_representation(observations, representation_function, batch_size):
# def discrete_mutual_info(mus, ys):
# def discrete_entropy(ys):
# def make_discretizer(target, num_bins=gin.REQUIRED,
#                      discretizer_fn=gin.REQUIRED):
# def _histogram_discretize(target, num_bins=gin.REQUIRED):
# def normalize_data(data, mean=None, stddev=None):
# def make_predictor_fn(predictor_fn=gin.REQUIRED):
# def logistic_regression_cv():
# def gradient_boosting_classifier():
# 
# 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/dendrogram.py
# def dendrogram_plot(matrix, output_dir, factor_names):
# def report_merges(z, num_factors):
# def _find(nodes, i):
# def _union(nodes, idx, idy, val, z, cluster_id, size, matrix, n_clusters,
#            idx_found):
# 
# Path: disentanglement_lib/visualize/visualize_scores.py
# def heat_square(matrix, output_dir, name, xlabel, ylabel, max_val=None,
#                 factor_names=None):
# def plot_matrix_squares(matrix, max_val, palette, size_scale, ax):
#   def to_color(val):
# def plot_bar_palette(palette, max_val, ax):
# def plot_recovery_vs_independent(matrix, output_dir, name):
# def precision(matrix, th):
# def recall(matrix, th):
# def bfs(matrix, to_visit, factors, codes, size):


import os
import numpy as np
import gin.tf
from absl import logging
from disentanglement_lib.evaluation.metrics import dci
from disentanglement_lib.evaluation.metrics import modularity_explicitness
from disentanglement_lib.evaluation.metrics import sap_score
from disentanglement_lib.evaluation.metrics import utils
from disentanglement_lib.utils import results
from disentanglement_lib.visualize import dendrogram
from disentanglement_lib.visualize import visualize_scores
    Unified scores.
  """
  mus = utils.obtain_representation(observations, representation_function,
                                    batch_size)
  assert labels.shape[1] == observations.shape[0], "Wrong labels shape."
  assert mus.shape[1] == observations.shape[0], "Wrong representation shape."

  mus_train, mus_test = utils.split_train_test(
      mus,
      train_percentage)
  ys_train, ys_test = utils.split_train_test(
      labels,
      train_percentage)
  return unified_scores(mus_train, ys_train, mus_test, ys_test, matrix_fns)


def unified_scores(mus_train, ys_train, mus_test, ys_test, matrix_fns,
                   artifact_dir=None, factor_names=None):
  """Computes unified scores."""

  scores = {}
  kws = {}
  for matrix_fn in matrix_fns:
    # Matrix should have shape [num_codes, num_factors].
    matrix = matrix_fn(mus_train, ys_train, mus_test, ys_test)
    matrix_name = matrix_fn.__name__
    if artifact_dir is not None:
      visualize_scores.heat_square(matrix.copy(), artifact_dir, matrix_name,
                                   "Latent codes", "Factors of Variation",
                                   factor_names=factor_names)

# Path: disentanglement_lib/evaluation/metrics/dci.py
# def compute_dci(ground_truth_data, representation_function, random_state,
#                 artifact_dir=None,
#                 num_train=gin.REQUIRED,
#                 num_test=gin.REQUIRED,
#                 batch_size=16):
# def _compute_dci(mus_train, ys_train, mus_test, ys_test):
# def compute_dci_on_fixed_data(observations, labels, representation_function,
#                               train_percentage=gin.REQUIRED, batch_size=100):
# def compute_importance_gbt(x_train, y_train, x_test, y_test):
# def disentanglement_per_code(importance_matrix):
# def disentanglement(importance_matrix):
# def completeness_per_factor(importance_matrix):
# def completeness(importance_matrix):
# 
# Path: disentanglement_lib/evaluation/metrics/modularity_explicitness.py
# def compute_modularity_explicitness(ground_truth_data,
#                                     representation_function,
#                                     random_state,
#                                     artifact_dir=None,
#                                     num_train=gin.REQUIRED,
#                                     num_test=gin.REQUIRED,
#                                     batch_size=16):
# def explicitness_per_factor(mus_train, y_train, mus_test, y_test):
# def modularity(mutual_information):
# 
# Path: disentanglement_lib/evaluation/metrics/sap_score.py
# def compute_sap(ground_truth_data,
#                 representation_function,
#                 random_state,
#                 artifact_dir=None,
#                 num_train=gin.REQUIRED,
#                 num_test=gin.REQUIRED,
#                 batch_size=16,
#                 continuous_factors=gin.REQUIRED):
# def _compute_sap(mus, ys, mus_test, ys_test, continuous_factors):
# def compute_sap_on_fixed_data(observations, labels, representation_function,
#                               train_percentage=gin.REQUIRED,
#                               continuous_factors=gin.REQUIRED,
#                               batch_size=100):
# def compute_score_matrix(mus, ys, mus_test, ys_test, continuous_factors):
# def compute_avg_diff_top_two(matrix):
# 
# Path: disentanglement_lib/evaluation/metrics/utils.py
# def generate_batch_factor_code(ground_truth_data, representation_function,
#                                num_points, random_state, batch_size):
# def split_train_test(observations, train_percentage):
# def obtain_representation(observations, representation_function, batch_size):
# def discrete_mutual_info(mus, ys):
# def discrete_entropy(ys):
# def make_discretizer(target, num_bins=gin.REQUIRED,
#                      discretizer_fn=gin.REQUIRED):
# def _histogram_discretize(target, num_bins=gin.REQUIRED):
# def normalize_data(data, mean=None, stddev=None):
# def make_predictor_fn(predictor_fn=gin.REQUIRED):
# def logistic_regression_cv():
# def gradient_boosting_classifier():
# 
# 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/dendrogram.py
# def dendrogram_plot(matrix, output_dir, factor_names):
# def report_merges(z, num_factors):
# def _find(nodes, i):
# def _union(nodes, idx, idy, val, z, cluster_id, size, matrix, n_clusters,
#            idx_found):
# 
# Path: disentanglement_lib/visualize/visualize_scores.py
# def heat_square(matrix, output_dir, name, xlabel, ylabel, max_val=None,
#                 factor_names=None):
# def plot_matrix_squares(matrix, max_val, palette, size_scale, ax):
#   def to_color(val):
# def plot_bar_palette(palette, max_val, ax):
# def plot_recovery_vs_independent(matrix, output_dir, name):
# def precision(matrix, th):
# def recall(matrix, th):
# def bfs(matrix, to_visit, factors, codes, size):


# Path: disentanglement_lib/evaluation/metrics/unified_scores.py
import os
import numpy as np
import gin.tf
from absl import logging
from disentanglement_lib.evaluation.metrics import dci
from disentanglement_lib.evaluation.metrics import modularity_explicitness
from disentanglement_lib.evaluation.metrics import sap_score
from disentanglement_lib.evaluation.metrics import utils
from disentanglement_lib.utils import results
from disentanglement_lib.visualize import dendrogram
from disentanglement_lib.visualize import visualize_scores

    Unified scores.
  """
  mus = utils.obtain_representation(observations, representation_function,
                                    batch_size)
  assert labels.shape[1] == observations.shape[0], "Wrong labels shape."
  assert mus.shape[1] == observations.shape[0], "Wrong representation shape."

  mus_train, mus_test = utils.split_train_test(
      mus,
      train_percentage)
  ys_train, ys_test = utils.split_train_test(
      labels,
      train_percentage)
  return unified_scores(mus_train, ys_train, mus_test, ys_test, matrix_fns)


def unified_scores(mus_train, ys_train, mus_test, ys_test, matrix_fns,
                   artifact_dir=None, factor_names=None):
  """Computes unified scores."""

  scores = {}
  kws = {}
  for matrix_fn in matrix_fns:
    # Matrix should have shape [num_codes, num_factors].
    matrix = matrix_fn(mus_train, ys_train, mus_test, ys_test)
    matrix_name = matrix_fn.__name__
    if artifact_dir is not None:
      visualize_scores.heat_square(matrix.copy(), artifact_dir, matrix_name,
                                   "Latent codes", "Factors of Variation",
                                   factor_names=factor_names)

      visualize_scores.plot_recovery_vs_independent(matrix.copy().T,