import collections


def reverse_map(
    applicable_label_dict,
    label_vocab = None):
  """Flip parenthood dict to map parents to children.

  Args:
    applicable_label_dict: e.g. output of get_applicable_label_dict.
    label_vocab: e.g. output of inference_lib.vocab_from_model_base_path

  Returns:
    collections.defaultdict of k, v where:
    k: originally the values in applicable_label_dict
    v: originally the keys in applicable_label_dict.
    The defaultdict returns an empty frozenset for keys that are not found.
    This behavior is desirable for lifted clan label normalizers, where
    keys may not imply themselves.
  """
  # This is technically the entire transitive closure, so it is safe for DAGs
  # (e.g. GO labels).

  children = collections.defaultdict(set)
  for child, parents in applicable_label_dict.items():
    # Avoid adding children which don't appear in the vocab.
    if label_vocab is None or child in label_vocab:
      for parent in parents:
        children[parent].add(child)
  children = {k: frozenset(v) for k, v in children.items()}
  return collections.defaultdict(frozenset, children.items())