import tempfile
import logging


def get_merged_vocabulary_file(vocabularies, padded_token=None):
  """Merges several vocabulary files into one temporary file.

  The TF object that loads the embedding expects a vocabulary file, to know
  which embeddings it should load.
  See tf.contrib.embedding.load_embedding_initializer.

  When we want to train/test on several datasets simultaneously we need to merge
  their vocabulary files into a single file.

  Args:
    vocabularies: Iterable of vocabularies. Each vocabulary should be
        a list of tokens.
    padded_token: If not None, add the padded_token to the first index.
  Returns:
    outfilename: Name of the merged file. Contains the union of all tokens in
        filenames, without duplicates, one token per line.
    vocabulary_size: Count of tokens in the merged file.
  """
  uniques = [set(vocabulary) for vocabulary in vocabularies]
  unique_merged = frozenset().union(*uniques)
  unique_merged_sorted = sorted(unique_merged)
  if padded_token is not None:
    # Add padded token as 0 index.
    unique_merged_sorted = [padded_token] + unique_merged_sorted
  vocabulary_size = len(unique_merged_sorted)
  outfile = tempfile.NamedTemporaryFile(delete=False)
  outfile.write(b'\n'.join(unique_merged_sorted))
  outfilename = outfile.name
  logging.info('Merged vocabulary file with %d tokens: %s', vocabulary_size,
               outfilename)
  outfile.close()
  return outfilename, vocabulary_size