#!/usr/bin/python
import codecs, pymatrices, time, sys, argparse


def generate_kbest_xml(outfile, simfile, matrixdir, matrixname, k, minnz, remove_prefix = None):
	matrix = pymatrices.LoadSuperMatrix(matrixdir, matrixname, False, False)
	features = matrix.GetRowFeatures()
	
	file_in = open(simfile)
	file_out = open(outfile, "w")
	line_no = 0
	rows = 0
	row_no = 0
	firstline = True
	file_out.write('<?xml version="1.0" encoding="UTF-8"?>\n')
	file_out.write('<thesaurus>\n')
	file_out.write('\t<description>\n')
	file_out.write('\t\t<generator>kbest.py - from similarity matrix %s </generator>\n' % matrixname)
	file_out.write('\t\t<date>%s</date>\n' %  time.strftime('%Y/%m/%d %H:%M'))
	file_out.write('\t</description>\n')
	file_out.write('\t<allsynsets>\n')
	for line in file_in:
		if firstline:
			firstline = False
			tokens = line.strip().split(" ")
			rows = int(tokens[0])
			continue

		line = line.strip()
		tokens = line.split(" ")

		if len(tokens) % 2 != 0 and len(line) > 0:
			print >> sys.stderr, "len(tokens) % 2 != 0, line:",line
			continue

		word = features[row_no].GetLabel()
		if remove_prefix and word.startswith(remove_prefix):
			word = word[len(remove_prefix):]

		if features[row_no].GetNZ() >= minnz:
			file_out.write('\t\t<synset lexicalunit="%s">\n' % word)
			vals = {}
			for tok_no in xrange(0, len(tokens), 2):
				id = int(tokens[tok_no])
				val = float(tokens[tok_no+1])
				if val in vals:
					vals[val].append(id)
				else:
					vals[val] = [id]

			added = 0
			keys = sorted( vals.keys() )
			for i in xrange(1, len(vals) +1):
				id = keys[ -i ]
				for word_id in vals[id]:
					if features[word_id].GetNZ() < minnz:
						continue
					word2 = features[word_id].GetLabel()
					if remove_prefix and word2.startswith(remove_prefix):
						word2 = word2[len(remove_prefix):]
					line = '\t\t\t<lu distance="%s">%s</lu>\n' % (id, word2)
					file_out.write( line )
					added += 1
					if added >= k: break
				if added >= k: break
			file_out.write('\t\t</synset>\n')
		
		line_no +=1
		row_no += 1

		if line_no % 100 == 0:
			print "Reading row %s / %s" % (row_no, rows )
	file_out.write('\t</allsynsets>\n')
	file_out.write('</thesaurus>\n')
	
	file_out.close()
def generate_kbest_plain(outfile, simfile, matrixdir, matrixname, k, minnz):
	matrix = pymatrices.LoadSuperMatrix(matrixdir, matrixname, False, False)
	features = matrix.GetRowFeatures()
	
	file_in = open(simfile)
	file_out = open(outfile, "w")
	line_no = 0
	rows = 0
	row_no = 0
	firstline = True
	for line in file_in:
		if firstline:
			firstline = False
			tokens = line.strip().split(" ")
			rows = int(tokens[0])
			continue

		line = line.strip()
		tokens = line.split(" ")

		if len(tokens) % 2 != 0 and len(line) > 0:
			print >> sys.stderr, "len(tokens) % 2 != 0, line:",line
			continue

		word = features[row_no].GetLabel()
		if features[row_no].GetNZ() >= minnz:
			file_out.write(word+"\n")
			vals = {}
			for tok_no in xrange(0, len(tokens), 2):
				id = int(tokens[tok_no])
				val = float(tokens[tok_no+1])
				if val in vals:
					vals[val].append(id)
				else:
					vals[val] = [id]

			added = 0
			keys = sorted( vals.keys() )
			for i in xrange(1, len(vals) +1):
				id = keys[ -i ]
				for word_id in vals[id]:
					if features[word_id].GetNZ() < minnz:
						continue
					word2 = features[word_id].GetLabel()
					line = "\t%s       %s\n" % (id, word2)
					file_out.write( line )
					added += 1
					if added >= k: break
				if added >= k: break
		else:
			file_out.write(word + " za malo NZ\n")
		
		line_no +=1
		row_no += 1

		if line_no % 100 == 0:
			print "Reading row %s / %s" % (row_no, rows )
	
	file_out.close()

def generate_kbest_pairs(outfile, simfile, matrixdir, matrixname, k, minnz):
	matrix = pymatrices.LoadSuperMatrix(matrixdir, matrixname, False, False)
	features = matrix.GetRowFeatures()
	
	file_in = open(simfile)
	file_out = open(outfile, "w")
	line_no = 0
	rows = 0
	row_no = 0
	firstline = True
	for line in file_in:
		if firstline:
			firstline = False
			tokens = line.strip().split(" ")
			rows = int(tokens[0])
			continue

		line = line.strip()
		tokens = line.split(" ")

		if len(tokens) % 2 != 0 and len(line) > 0:
			print >> sys.stderr, "len(tokens) % 2 != 0, line:",line
			continue

		word = features[row_no].GetLabel()
		if features[row_no].GetNZ() >= minnz:
			vals = {}
			for tok_no in xrange(0, len(tokens), 2):
				id = int(tokens[tok_no])
				val = float(tokens[tok_no+1])
				if val in vals:
					vals[val].append(id)
				else:
					vals[val] = [id]

			added = 0
			keys = sorted( vals.keys() )
			for i in xrange(1, len(vals) +1):
				id = keys[ -i ]
				for word_id in vals[id]:
					if features[word_id].GetNZ() < minnz:
						continue
					word2 = features[word_id].GetLabel()
					line = "%s; %s; %f\n" % (word.replace("subst:", ""),word2.replace("subst:", ""), id)
					file_out.write( line )
					added += 1
					if added >= k: break
				if added >= k: break
		else:
			print (word + " za malo NZ")
		
		line_no +=1
		row_no += 1

		if line_no % 100 == 0:
			print "Reading row %s / %s" % (row_no, rows )
	
	file_out.close()	
	
if __name__ == '__main__':
	parser = argparse.ArgumentParser(description='kbest')
	parser.add_argument('--matrix-dir', action="store", help="Path to matrix directory")
	parser.add_argument('--matrix-name', action="store", help="Matrix name")
	parser.add_argument('--sim-file', action="store", help="Path to similarity file")
	parser.add_argument('--out-file', action="store", help="Path to output file")
	parser.add_argument('--k', action="store", help="Value of k")
	parser.add_argument('--minnz', action="store", help="Minimum non-zero value ")
	parser.add_argument('--flex', action="store", default="subst:", help="Flex type")	
	args = parser.parse_args()
	
	print args

	#filename = "kipi_knigi_rzepa_verb_3000_3010_3030_all_imiemodsubst_rwf_zscore_nz5.txt"
	#generate_kbest_plain(filename, simfile, matrixdir, matrixname, k, minnz)

	generate_kbest_xml(args.out_file, args.sim_file, args.matrix_dir, args.matrix_name, args.k, args.minnz, args.flex)
