import json


def convert_neighbor_to_json(infn, outfn):
  """convert data format for nearest neighbors.

  Args:
    infn: tsv file downloaded from MSMARCO dataset.
    outfn: path to output neighors.

  Returns:
    The output neighbor data format is as follows:
    dict(qID=[list of ranked neighbors in the format of [pID, score]])
  """
  data = {}
  fns = infn.split(',')
  for fn in fns:
    for line in open(fn):
      elements = line.split('\t')
      if elements[0] not in data:
        data[elements[0]] = []
      # no BM25 score provided in MSMARCO dataset, use 0 instead
      neighbor = [elements[1], 0]  # replace 0 with BM25 score
      # data[elements[0]] needs to be sorted by BM25 score
      data[elements[0]].append(neighbor)

  with open(outfn, 'w') as outfile:
    json.dump(data, outfile)
  print('Processed: ' + infn)
  print('Output to: ' + outfn)
  print('Number of entries: ' + str(len(data)))
  return data