# Path: meta_dataset/data/imagenet_specification.py
# FLAGS = tf.flags.FLAGS
# class Synset(object):
#   def __init__(self, wn_id, words, children, parents):
# def get_node_ancestors(synset):
# def get_ancestors(synsets):
# def isolate_graph(nodes):
# def isolate_node(node):
# def collapse(nodes):
#   def collapse_once(nodes):
# def get_leaves(nodes):
# def get_synsets_from_ids(wn_ids, synsets):
# def get_spanning_leaves(nodes):
# def get_num_spanning_images(spanning_leaves, num_leaf_images):
# def create_sampling_graph(synsets, root=None):
# def propose_valid_test_roots(spanning_leaves,
#                              margin=50,
#                              desired_num_valid_classes=150,
#                              desired_num_test_classes=150):
#   def _sort_key(synset_and_leaves):
# def get_class_splits(spanning_leaves, valid_test_roots=None, **kwargs):
# def init_split_subgraphs(class_splits, spanning_leaves, valid_test_roots):
# def copy_graph(nodes, root_wn_id=None):
# def create_splits(spanning_leaves, split_enum, valid_test_roots=None, **kwargs):
# def get_synset_by_wnid(wnid, graph):
# def is_descendent(d, a):
# def get_upward_paths_from(start, end=None):
#   def is_end_node(n):
# def find_lowest_common_in_paths(path_a, path_b):
# def get_lowest_common_ancestor(leaf_a, leaf_b, path='longest'):
# def get_num_synset_2012_images(path, synsets_2012, files_to_skip=None):
# def export_graph(nodes):
# def import_graph(node_representations):
# def create_imagenet_specification(split_enum,
#                                   files_to_skip,
#                                   path_to_num_leaf_images=None,
#                                   train_split_only=False,
#                                   log_stats=True):
# 
# Path: meta_dataset/data/learning_spec.py
# class Split(enum.Enum):
# class BatchSpecification(
#     collections.namedtuple('BatchSpecification', 'split, batch_size')):
# class EpisodeSpecification(
#     collections.namedtuple(
#         'EpisodeSpecification',
#         'split, num_classes, num_train_examples, num_test_examples')):
#   TRAIN = 0
#   VALID = 1
#   TEST = 2


from meta_dataset.data import imagenet_specification as imagenet_spec
from meta_dataset.data import learning_spec
from six.moves import range
import numpy as np
import tensorflow.compat.v1 as tf

  def test_toy_root_proposer(self):
    specification = create_toy_graph()
    _, toy_span_leaves, _ = specification
    valid_test_roots = imagenet_spec.propose_valid_test_roots(
        toy_span_leaves,
        margin=TOY_MARGIN,
        desired_num_valid_classes=DESIRED_TOY_NUM_VALID_CLASSES,
        desired_num_test_classes=DESIRED_TOY_NUM_TEST_CLASSES)
    self.validate_roots(valid_test_roots, toy_span_leaves)

    # returns the lists of id's of classes belonging to each split
    # unlike create_splits which returns the nodes of the three subgraphs that
    # are constructed for the different splits
    splits, _ = imagenet_spec.get_class_splits(
        toy_span_leaves, valid_test_roots=valid_test_roots)
    self.validate_splits(splits, toy_span_leaves)


class ImagenetSpecificationTest(tf.test.TestCase):

  def validate_num_span_images(self, span_leaves, num_span_images):
    # Ensure that the number of images spanned by each node is exactly the
    # number of images living in the leaves spanned by that node
    for node, leaves in span_leaves.items():
      self.assertEqual(num_span_images[node],
                       sum([num_span_images[l] for l in leaves]))

  def validate_splits(self, splits):
    """Check the correctness of the class splits."""

# Path: meta_dataset/data/imagenet_specification.py
# FLAGS = tf.flags.FLAGS
# class Synset(object):
#   def __init__(self, wn_id, words, children, parents):
# def get_node_ancestors(synset):
# def get_ancestors(synsets):
# def isolate_graph(nodes):
# def isolate_node(node):
# def collapse(nodes):
#   def collapse_once(nodes):
# def get_leaves(nodes):
# def get_synsets_from_ids(wn_ids, synsets):
# def get_spanning_leaves(nodes):
# def get_num_spanning_images(spanning_leaves, num_leaf_images):
# def create_sampling_graph(synsets, root=None):
# def propose_valid_test_roots(spanning_leaves,
#                              margin=50,
#                              desired_num_valid_classes=150,
#                              desired_num_test_classes=150):
#   def _sort_key(synset_and_leaves):
# def get_class_splits(spanning_leaves, valid_test_roots=None, **kwargs):
# def init_split_subgraphs(class_splits, spanning_leaves, valid_test_roots):
# def copy_graph(nodes, root_wn_id=None):
# def create_splits(spanning_leaves, split_enum, valid_test_roots=None, **kwargs):
# def get_synset_by_wnid(wnid, graph):
# def is_descendent(d, a):
# def get_upward_paths_from(start, end=None):
#   def is_end_node(n):
# def find_lowest_common_in_paths(path_a, path_b):
# def get_lowest_common_ancestor(leaf_a, leaf_b, path='longest'):
# def get_num_synset_2012_images(path, synsets_2012, files_to_skip=None):
# def export_graph(nodes):
# def import_graph(node_representations):
# def create_imagenet_specification(split_enum,
#                                   files_to_skip,
#                                   path_to_num_leaf_images=None,
#                                   train_split_only=False,
#                                   log_stats=True):
# 
# Path: meta_dataset/data/learning_spec.py
# class Split(enum.Enum):
# class BatchSpecification(
#     collections.namedtuple('BatchSpecification', 'split, batch_size')):
# class EpisodeSpecification(
#     collections.namedtuple(
#         'EpisodeSpecification',
#         'split, num_classes, num_train_examples, num_test_examples')):
#   TRAIN = 0
#   VALID = 1
#   TEST = 2


# Path: meta_dataset/data/imagenet_specification_test.py
from meta_dataset.data import imagenet_specification as imagenet_spec
from meta_dataset.data import learning_spec
from six.moves import range
import numpy as np
import tensorflow.compat.v1 as tf


  def test_toy_root_proposer(self):
    specification = create_toy_graph()
    _, toy_span_leaves, _ = specification
    valid_test_roots = imagenet_spec.propose_valid_test_roots(
        toy_span_leaves,
        margin=TOY_MARGIN,
        desired_num_valid_classes=DESIRED_TOY_NUM_VALID_CLASSES,
        desired_num_test_classes=DESIRED_TOY_NUM_TEST_CLASSES)
    self.validate_roots(valid_test_roots, toy_span_leaves)

    # returns the lists of id's of classes belonging to each split
    # unlike create_splits which returns the nodes of the three subgraphs that
    # are constructed for the different splits
    splits, _ = imagenet_spec.get_class_splits(
        toy_span_leaves, valid_test_roots=valid_test_roots)
    self.validate_splits(splits, toy_span_leaves)


class ImagenetSpecificationTest(tf.test.TestCase):

  def validate_num_span_images(self, span_leaves, num_span_images):
    # Ensure that the number of images spanned by each node is exactly the
    # number of images living in the leaves spanned by that node
    for node, leaves in span_leaves.items():
      self.assertEqual(num_span_images[node],
                       sum([num_span_images[l] for l in leaves]))

  def validate_splits(self, splits):
    """Check the correctness of the class splits."""

    train_graph = splits[learning_spec.Split.TRAIN]