import collections


def _combine_gref_anns(gref_anns):
  """Reorganize G-Ref by matching annotations to image_id.

  Also map referring expressions to annotation id.
  Args:
    gref_anns: a dict of gref_anns annotations.

  Returns:
    combine_anns: a list of dict based on gref_anns['images'].
      Each dict has a key 'instances' corresponding to a list of annotations
      of that image.
  """
  img_id2annotations = collections.defaultdict(list)
  for _, v in gref_anns['annotations'].items():
    # Add ref expression to annotation.
    # this adds the referent, the raw string, tokens, and the referring
    # expression id to the annotation.
    ref_list = []
    for rexp_id in sorted(v['refexp_ids']):
      r = gref_anns['refexps'][str(rexp_id)]
      if 'word' in r['parse']['referent']:
        r['referent'] = r['parse']['referent']['word']
      else:
        r['referent'] = ''
      r.pop('parse')
      ref_list.append(r)

    v['refexps'] = ref_list
    img_id2annotations[int(v['image_id'])].append(v)

  img_id2annotations = {
      k: list(sorted(v, key=lambda a: a['annotation_id']))
      for k, v in img_id2annotations.items()
  }
  combined_anns = []
  for _, v in gref_anns['images'].items():
    v['instances'] = img_id2annotations[v['image_id']]
    combined_anns.append((v['image_id'], v))

  combined_anns = sorted(combined_anns, key=lambda a: a[0])
  return combined_anns