# 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/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():


from disentanglement_lib.evaluation.metrics import dci
from disentanglement_lib.evaluation.metrics import utils
from six.moves import range
import numpy as np
import gin.tf
  Returns:
    Tuple with reduced representations for the training and test set.
  """
  importance_matrix = correlation_measure(mus_train, ys_train, mus_test,
                                          ys_test)
  factor_of_interest_importance = importance_matrix[:, factor_of_interest]
  factor_to_remove_index = np.argmax(factor_of_interest_importance)
  # Remove the factor of variation above from the representation
  reduced_representation_train = np.delete(
      mus_train.copy(), factor_to_remove_index, axis=0)
  reduced_representation_test = np.delete(
      mus_test.copy(), factor_to_remove_index, axis=0)
  return reduced_representation_train, reduced_representation_test


@gin.configurable(
    "factorwise_dci",
    blacklist=["mus_train", "ys_train", "mus_test", "ys_test"])
def compute_factorwise_dci(mus_train, ys_train, mus_test, ys_test):
  """Computes the DCI importance matrix of the attributes.

  Args:
    mus_train: latent means of the training batch.
    ys_train: labels of the training batch.
    mus_test: latent means of the test batch.
    ys_test: labels of the test batch.

  Returns:
    Matrix with importance scores.
  """

# 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/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/evaluation/metrics/reduced_downstream_task.py
from disentanglement_lib.evaluation.metrics import dci
from disentanglement_lib.evaluation.metrics import utils
from six.moves import range
import numpy as np
import gin.tf

  Returns:
    Tuple with reduced representations for the training and test set.
  """
  importance_matrix = correlation_measure(mus_train, ys_train, mus_test,
                                          ys_test)
  factor_of_interest_importance = importance_matrix[:, factor_of_interest]
  factor_to_remove_index = np.argmax(factor_of_interest_importance)
  # Remove the factor of variation above from the representation
  reduced_representation_train = np.delete(
      mus_train.copy(), factor_to_remove_index, axis=0)
  reduced_representation_test = np.delete(
      mus_test.copy(), factor_to_remove_index, axis=0)
  return reduced_representation_train, reduced_representation_test


@gin.configurable(
    "factorwise_dci",
    blacklist=["mus_train", "ys_train", "mus_test", "ys_test"])
def compute_factorwise_dci(mus_train, ys_train, mus_test, ys_test):
  """Computes the DCI importance matrix of the attributes.

  Args:
    mus_train: latent means of the training batch.
    ys_train: labels of the training batch.
    mus_test: latent means of the test batch.
    ys_test: labels of the test batch.

  Returns:
    Matrix with importance scores.
  """

  importance_matrix, _, _ = dci.compute_importance_gbt(mus_train, ys_train,