#!/bin/python

'''
Simulating a TF selector on redundantly large k-best lists
Created on 2011-04-18

@author: Jacek Bochenek
'''


import codecs
import os, sys

def getKBest(kbest_file):
	kbest = {}
	current = ""
	for line in codecs.open(kbest_file,"r","utf-8"):
		if not line[0] == "\t":
			kbest[line.strip()] = []
			current = line.strip()
		else:
			line = line.strip()
			contents = line.split(";")
			kbest[current].append((contents[0],contents[1]))
	return kbest

def getRowFeatures(featuresFile):
	nz = {}
	tf = {}
	TF = {}
	ent = {}
	firstLine = True
	for line in codecs.open(featuresFile,"r","utf-8"):
		if firstLine:
			firstLine = False
			continue
		line = line.strip()
		contents = line.split(";")
		word = contents[0]
		nz[word] = int(contents[1])
		tf[word] = int(contents[2])
		TF[word] = int(contents[3])
		ent[word] = float(contents[4])
	return (nz,tf,TF,ent)

def filterKBest(kbest_file,features_file,min_tf,target_k,out_file):
	ofile = codecs.open(out_file,"w","utf-8")
	kbest = getKBest(kbest_file)
	features = getRowFeatures(features_file)
	for key in sorted(kbest.keys()):
		if features[1][key] > min_tf and key[:6] == "subst:":
			ofile.write(key+"\n")
		else:
			continue
		current = 0
		for i in xrange(0,len(kbest[key])):
			if current == target_k:
				break
			if features[1][kbest[key][i][0]] > min_tf and "-" not in kbest[key][i][0]:#  and kbest[key][i][0][:4] == "fin:":
				ofile.write("\t" + kbest[key][i][0] + ";" + kbest[key][i][1] + "\n")
				current += 1
	return kbest

def main():
	if len(sys.argv) < 6:
		print "Usage: " + sys.argv[0] + " kbest_file features_file min_tf target_k out_file"
		sys.exit(-1)
	kbest_file = sys.argv[1]
	features_file = sys.argv[2]
	min_tf = int(sys.argv[3])
	target_k = int(sys.argv[4])
	out_file = sys.argv[5]
	filterKBest(kbest_file,features_file,min_tf,target_k,out_file)

if __name__ == '__main__':
	main()
	

