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


from absl.testing import absltest
from absl.testing import parameterized
from disentanglement_lib.visualize import visualize_scores
import numpy as np
                 [0, 0, 0, 0, 0],
                 [0, 0, 0, 0, 0],
                 [0, 0, 0, 0, 0],
                 [0, 0, 0, 1, 1]]),
       [(0, 0)], [1, 0, 0, 0, 0], [1, 1, 0, 0, 0], 3),
      (np.array([[1, 1, 0, 0, 0],
                 [0, 0, 0, 0, 0],
                 [0, 0, 0, 0, 0],
                 [0, 1, 0, 0, 0],
                 [0, 0, 0, 1, 1]]),
       [(0, 0)], [1, 0, 0, 1, 0], [1, 1, 0, 0, 0], 4)
  )
  def test_bfs(self, matrix, to_visit, visited_factors, visited_codes,
               visited_size):
    factors = np.zeros(matrix.shape[0])
    codes = np.zeros(matrix.shape[1])
    factors, codes, size = visualize_scores.bfs(
        matrix, to_visit, factors, codes, 1)

    self.assertEqual((factors == visited_factors).all(), True)
    self.assertEqual((codes == visited_codes).all(), True)
    self.assertEqual(size, visited_size)

  @parameterized.parameters(
      (np.zeros((5, 5)), 0),
      (np.eye(5), 5),
      (np.triu(np.ones(5)), 1),
      (np.ones((5, 5)), 1)
  )
  def test_precision(self, matrix, cc):

# 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/visualize/visualize_scores_test.py
from absl.testing import absltest
from absl.testing import parameterized
from disentanglement_lib.visualize import visualize_scores
import numpy as np

                 [0, 0, 0, 0, 0],
                 [0, 0, 0, 0, 0],
                 [0, 0, 0, 0, 0],
                 [0, 0, 0, 1, 1]]),
       [(0, 0)], [1, 0, 0, 0, 0], [1, 1, 0, 0, 0], 3),
      (np.array([[1, 1, 0, 0, 0],
                 [0, 0, 0, 0, 0],
                 [0, 0, 0, 0, 0],
                 [0, 1, 0, 0, 0],
                 [0, 0, 0, 1, 1]]),
       [(0, 0)], [1, 0, 0, 1, 0], [1, 1, 0, 0, 0], 4)
  )
  def test_bfs(self, matrix, to_visit, visited_factors, visited_codes,
               visited_size):
    factors = np.zeros(matrix.shape[0])
    codes = np.zeros(matrix.shape[1])
    factors, codes, size = visualize_scores.bfs(
        matrix, to_visit, factors, codes, 1)

    self.assertEqual((factors == visited_factors).all(), True)
    self.assertEqual((codes == visited_codes).all(), True)
    self.assertEqual(size, visited_size)

  @parameterized.parameters(
      (np.zeros((5, 5)), 0),
      (np.eye(5), 5),
      (np.triu(np.ones(5)), 1),
      (np.ones((5, 5)), 1)
  )
  def test_precision(self, matrix, cc):

    cc_count = visualize_scores.precision(matrix, 1)