import codecs
import sys
import pymatrices
import random
import os.path


if len(sys.argv) != 9:
	print "Usage: " + sys.argv[0] + " <wn_words> <matrix_flex> <matrix_dir> <matrix_name> <rare_threshold> <frequent_threshold> <sample_size> <output_dir>"
	sys.exit(1)

wn = {}
for line in codecs.open(sys.argv[1],"r","utf-8"):
	line = line.encode('utf-8')
	wn[line.strip()] = 1

matrix_flex = sys.argv[2]
matrix_dir = sys.argv[3]
matrix_name = sys.argv[4]
rare_threshold = int(sys.argv[5])
frequent_threshold = int(sys.argv[6])
sample_size = int(sys.argv[7])
output_dir = sys.argv[8]

def remove_flex(flex,word):
	return word[len(flex):]

matrix = pymatrices.LoadSuperMatrix(matrix_dir, matrix_name, False, False)
rowfeatures = matrix.GetRowFeatures()

all_words = []
frequent_words = []
rare_words = []

for word in wn:
	label = matrix_flex + ":" + word
	#label = label.encode('utf-8')
	#print type(word)
	word_id = matrix.GetRowIndex(label)
	if word_id > len(rowfeatures):
		#print "$" + label + "$ not in matrix!"
		#print word_id
		continue
	all_words.append(word)
	tf = rowfeatures[word_id].GetGlobalTermFrequency()
	if tf == 0:
		continue
	if tf >= frequent_threshold:
		frequent_words.append(word)
	if tf <= rare_threshold:
		rare_words.append(word)

random.seed(0)

random.shuffle(all_words)
random.shuffle(frequent_words)
random.shuffle(rare_words)

rarefile = codecs.open(os.path.join(output_dir,"rare.txt"),"w","utf-8")
allfile = codecs.open(os.path.join(output_dir,"all.txt"),"w","utf-8")
frequentfile = codecs.open(os.path.join(output_dir,"frequent.txt"),"w","utf-8")

for i in xrange(0,sample_size):
	if i < len(rare_words):
		rarefile.write(rare_words[i].decode('utf-8') + "\n")
	if i < len(all_words):
		allfile.write(all_words[i].decode('utf-8') + "\n")
	if i < len(frequent_words):
		frequentfile.write(frequent_words[i].decode('utf-8') + "\n")
	
	
	
