# Path: tensorflow_transform/info_theory.py
# def calculate_partial_expected_mutual_information(n, x_i, y_j):
# def calculate_partial_mutual_information(n_ij, x_i, y_j, n):
# def _hypergeometric_pmf(n, x_i, y_j):
# def _logfactorial(n):
# 
# Path: tensorflow_transform/test_case.py
# def is_tf_api_version_1():
# def cross_named_parameters(*args):
#   def _cross_test_cases(parameters_list):
# def parameters(*testcases):
#   def wrapper(fn):
#     def to_arg_dict(testcase):
# def cross_parameters(*args):
# def _make_placeholder(tensor_spec):
# def _graph_function_handler(input_signature):
#   def wrapper(fn):
#     def _run_graph(*inputs):
# def _ragged_value_as_constant(value, dtype):
# def _wrap_as_constant(value, tensor_spec):
# def _eager_function_handler(input_signature):
#   def wrapper(fn):
#     def _run_eagerly(*inputs):  # pylint: disable=missing-docstring
# def _tf_function_function_handler(input_signature):
#   def wrapper(fn):
# def is_external_environment():
# def skip_if_external_environment(reason):
# def skip_if_not_tf2(reason):
# def cross_with_function_handlers(parameters_list):
#   def assertDataCloseOrEqual(self, a_data, b_data):
#   def _assertValuesCloseOrEqual(self, a_value, b_value, msg=None):
#   def AssertVocabularyContents(self, vocab_file_path, file_contents):
#   def WriteRenderedDotFile(self, dot_string, output_file=None):
#   def _NumpyArraysToLists(self, maybe_arrays):
#   def _SortedDicts(self, list_of_dicts):
#   def _SortedData(self, list_of_dicts_of_arrays):
# FUNCTION_HANDLERS = [
#     dict(testcase_name='graph',
#          function_handler=_graph_function_handler),
#     dict(testcase_name='eager',
#          function_handler=_eager_function_handler),
#     dict(testcase_name='tf_function',
#          function_handler=_tf_function_function_handler)
# ]
# class TransformTestCase(parameterized.TestCase, tf.test.TestCase):


from tensorflow_transform import info_theory
from tensorflow_transform import test_case
import unittest

def _make_hypergeometric_pmf_sum_up_to_one_parameters():
  start = 1000
  end = 10000
  range_length = end - start
  num_chunks = 15
  assert range_length % num_chunks == 0
  chunk_size = int(range_length / num_chunks)
  sub_ranges = [(x, x + chunk_size) for x in range(start, end, chunk_size)]
  return [  # pylint: disable=g-complex-comprehension
      dict(
          testcase_name='{}_to_{}'.format(a, b),
          test_range=range(a, b),
          n=end,
          y_j=start) for a, b in sub_ranges
  ]


class InfoTheoryTest(test_case.TransformTestCase):

  def testHypergeometricPmf(self):
    expected_results = [(0, 0.75), (1, 0.25)]
    results = list(info_theory._hypergeometric_pmf(4, 1, 1))
    for expected_result, result in zip(expected_results, results):
      self.assertEqual(expected_result[0], result[0])
      self.assertNear(expected_result[1], result[1], EPSILON)

  def testHypergeometricPmf_LargeN(self):
    expected_results = [(0, 0.9508937), (1, 0.0482198), (2, 0.0008794),
                        (3, 7.1e-06), (4, 2.5e-08), (5, 0.0)]

# Path: tensorflow_transform/info_theory.py
# def calculate_partial_expected_mutual_information(n, x_i, y_j):
# def calculate_partial_mutual_information(n_ij, x_i, y_j, n):
# def _hypergeometric_pmf(n, x_i, y_j):
# def _logfactorial(n):
# 
# Path: tensorflow_transform/test_case.py
# def is_tf_api_version_1():
# def cross_named_parameters(*args):
#   def _cross_test_cases(parameters_list):
# def parameters(*testcases):
#   def wrapper(fn):
#     def to_arg_dict(testcase):
# def cross_parameters(*args):
# def _make_placeholder(tensor_spec):
# def _graph_function_handler(input_signature):
#   def wrapper(fn):
#     def _run_graph(*inputs):
# def _ragged_value_as_constant(value, dtype):
# def _wrap_as_constant(value, tensor_spec):
# def _eager_function_handler(input_signature):
#   def wrapper(fn):
#     def _run_eagerly(*inputs):  # pylint: disable=missing-docstring
# def _tf_function_function_handler(input_signature):
#   def wrapper(fn):
# def is_external_environment():
# def skip_if_external_environment(reason):
# def skip_if_not_tf2(reason):
# def cross_with_function_handlers(parameters_list):
#   def assertDataCloseOrEqual(self, a_data, b_data):
#   def _assertValuesCloseOrEqual(self, a_value, b_value, msg=None):
#   def AssertVocabularyContents(self, vocab_file_path, file_contents):
#   def WriteRenderedDotFile(self, dot_string, output_file=None):
#   def _NumpyArraysToLists(self, maybe_arrays):
#   def _SortedDicts(self, list_of_dicts):
#   def _SortedData(self, list_of_dicts_of_arrays):
# FUNCTION_HANDLERS = [
#     dict(testcase_name='graph',
#          function_handler=_graph_function_handler),
#     dict(testcase_name='eager',
#          function_handler=_eager_function_handler),
#     dict(testcase_name='tf_function',
#          function_handler=_tf_function_function_handler)
# ]
# class TransformTestCase(parameterized.TestCase, tf.test.TestCase):


# Path: tensorflow_transform/info_theory_test.py
from tensorflow_transform import info_theory
from tensorflow_transform import test_case
import unittest


def _make_hypergeometric_pmf_sum_up_to_one_parameters():
  start = 1000
  end = 10000
  range_length = end - start
  num_chunks = 15
  assert range_length % num_chunks == 0
  chunk_size = int(range_length / num_chunks)
  sub_ranges = [(x, x + chunk_size) for x in range(start, end, chunk_size)]
  return [  # pylint: disable=g-complex-comprehension
      dict(
          testcase_name='{}_to_{}'.format(a, b),
          test_range=range(a, b),
          n=end,
          y_j=start) for a, b in sub_ranges
  ]


class InfoTheoryTest(test_case.TransformTestCase):

  def testHypergeometricPmf(self):
    expected_results = [(0, 0.75), (1, 0.25)]
    results = list(info_theory._hypergeometric_pmf(4, 1, 1))
    for expected_result, result in zip(expected_results, results):
      self.assertEqual(expected_result[0], result[0])
      self.assertNear(expected_result[1], result[1], EPSILON)

  def testHypergeometricPmf_LargeN(self):
    expected_results = [(0, 0.9508937), (1, 0.0482198), (2, 0.0008794),
                        (3, 7.1e-06), (4, 2.5e-08), (5, 0.0)]

    results = list(info_theory._hypergeometric_pmf(1000, 5, 10))