# Path: disentanglement_lib/data/ground_truth/dummy_data.py
# class IdentityObservationsData(ground_truth_data.GroundTruthData):
# class DummyData(ground_truth_data.GroundTruthData):
#   def num_factors(self):
#   def observation_shape(self):
#   def factors_num_values(self):
#   def sample_factors(self, num, random_state):
#   def sample_observations_from_factors(self, factors, random_state):
#   def factor_names(self):
#   def num_factors(self):
#   def factors_num_values(self):
#   def observation_shape(self):
#   def sample_factors(self, num, random_state):
#   def sample_observations_from_factors(self, factors, random_state):
# 
# Path: disentanglement_lib/evaluation/udr/metrics/udr.py
# def relative_strength_disentanglement(corr_matrix):
# def spearman_correlation_conv(vec1, vec2):
# def lasso_correlation_matrix(vec1, vec2, random_state=None):
# def _generate_representation_batch(ground_truth_data, representation_functions,
#                                    batch_size, random_state):
# def _generate_representation_dataset(ground_truth_data,
#                                      representation_functions, batch_size,
#                                      num_data_points, random_state):
# def compute_udr_sklearn(ground_truth_data,
#                         representation_functions,
#                         random_state,
#                         batch_size,
#                         num_data_points,
#                         correlation_matrix="lasso",
#                         filter_low_kl=True,
#                         include_raw_correlations=True,
#                         kl_filter_threshold=0.01):


from absl.testing import absltest
from disentanglement_lib.data.ground_truth import dummy_data
from disentanglement_lib.evaluation.udr.metrics import udr
import numpy as np
    ground_truth_data = dummy_data.DummyData()
    random_state = np.random.RandomState(0)
    num_factors = ground_truth_data.num_factors
    batch_size = 10
    num_data_points = 1000

    permutation = np.random.permutation(num_factors)
    sign_inverse = np.random.choice(num_factors, int(num_factors / 2))

    def rep_fn1(data):
      return (np.reshape(data, (batch_size, -1))[:, :num_factors],
              np.ones(num_factors))

    # Should be invariant to permutation and sign inverse.
    def rep_fn2(data):
      raw_representation = np.reshape(data, (batch_size, -1))[:, :num_factors]
      perm_rep = raw_representation[:, permutation]
      perm_rep[:, sign_inverse] = -1.0 * perm_rep[:, sign_inverse]
      return perm_rep, np.ones(num_factors)

    scores = udr.compute_udr_sklearn(
        ground_truth_data, [rep_fn1, rep_fn2],
        random_state,
        batch_size,
        num_data_points,
        correlation_matrix="lasso")
    self.assertBetween(scores["model_scores"][0], 0.8, 1.0)
    self.assertBetween(scores["model_scores"][1], 0.8, 1.0)

  def test_metric_kl(self):

# Path: disentanglement_lib/data/ground_truth/dummy_data.py
# class IdentityObservationsData(ground_truth_data.GroundTruthData):
# class DummyData(ground_truth_data.GroundTruthData):
#   def num_factors(self):
#   def observation_shape(self):
#   def factors_num_values(self):
#   def sample_factors(self, num, random_state):
#   def sample_observations_from_factors(self, factors, random_state):
#   def factor_names(self):
#   def num_factors(self):
#   def factors_num_values(self):
#   def observation_shape(self):
#   def sample_factors(self, num, random_state):
#   def sample_observations_from_factors(self, factors, random_state):
# 
# Path: disentanglement_lib/evaluation/udr/metrics/udr.py
# def relative_strength_disentanglement(corr_matrix):
# def spearman_correlation_conv(vec1, vec2):
# def lasso_correlation_matrix(vec1, vec2, random_state=None):
# def _generate_representation_batch(ground_truth_data, representation_functions,
#                                    batch_size, random_state):
# def _generate_representation_dataset(ground_truth_data,
#                                      representation_functions, batch_size,
#                                      num_data_points, random_state):
# def compute_udr_sklearn(ground_truth_data,
#                         representation_functions,
#                         random_state,
#                         batch_size,
#                         num_data_points,
#                         correlation_matrix="lasso",
#                         filter_low_kl=True,
#                         include_raw_correlations=True,
#                         kl_filter_threshold=0.01):


# Path: disentanglement_lib/evaluation/udr/metrics/udr_test.py
from absl.testing import absltest
from disentanglement_lib.data.ground_truth import dummy_data
from disentanglement_lib.evaluation.udr.metrics import udr
import numpy as np

    ground_truth_data = dummy_data.DummyData()
    random_state = np.random.RandomState(0)
    num_factors = ground_truth_data.num_factors
    batch_size = 10
    num_data_points = 1000

    permutation = np.random.permutation(num_factors)
    sign_inverse = np.random.choice(num_factors, int(num_factors / 2))

    def rep_fn1(data):
      return (np.reshape(data, (batch_size, -1))[:, :num_factors],
              np.ones(num_factors))

    # Should be invariant to permutation and sign inverse.
    def rep_fn2(data):
      raw_representation = np.reshape(data, (batch_size, -1))[:, :num_factors]
      perm_rep = raw_representation[:, permutation]
      perm_rep[:, sign_inverse] = -1.0 * perm_rep[:, sign_inverse]
      return perm_rep, np.ones(num_factors)

    scores = udr.compute_udr_sklearn(
        ground_truth_data, [rep_fn1, rep_fn2],
        random_state,
        batch_size,
        num_data_points,
        correlation_matrix="lasso")
    self.assertBetween(scores["model_scores"][0], 0.8, 1.0)
    self.assertBetween(scores["model_scores"][1], 0.8, 1.0)

  def test_metric_kl(self):

    ground_truth_data = dummy_data.DummyData()