#!/usr/bin/python

import csv
import re
import argparse

parser = argparse.ArgumentParser(description='After annotators fill in empty TSV file generated by bibs-to-tsv.py, '
											 'this converts the filled-in TSV data into the format used for other TSV '
											 'gold data')
parser.add_argument('-i', metavar='INPUT.TSV', default='mentions-filled.tsv', help='Filename for filled-in TSV')
parser.add_argument('-o', metavar='OUTPUT.TSV', default='mentions.tsv', help='Filename for final gold TSV')
args = parser.parse_args()

with open(args.i) as mentions, open(args.o, 'w') as gold_writer:
	mentions = csv.reader(mentions, delimiter='\t')
	next(mentions) # skip header row
	gold_writer = csv.writer(gold_writer, delimiter='\t')
	gold = {}
	for paper, _, _, context, mention in mentions:
		if paper not in gold:
			gold[paper] = []
		cleaned_context = re.sub(r'\s+', ' ', context.strip())
		if not cleaned_context:
			continue
		gold[paper].append("{0}|{1}".format(cleaned_context, re.sub(r'[()]', '', mention)))
	for paper, bib_entries in gold.items():
		if len(bib_entries) > 0:
			gold_writer.writerow([paper] + bib_entries)