#!/usr/bin/python

import codecs, sys, numpy, pylab, random, pymatrices

def load_udict(filename, encoding = 'cp1250', flexfilter = set(['subst']), mintf=5):
	words = []
	sumtf = 0
	for line in codecs.open(filename, 'r', encoding):
		line = line.strip()
		flexwordcount = line.split(':')
		if len(flexwordcount) != 3:
			print >> sys.stderr, "Error in line:", line
			continue
		flex = flexwordcount[0].strip()
		if not flex in flexfilter:
			continue
		word = flexwordcount[1].strip()
		count = int(flexwordcount[2][1:].strip())
		if count < mintf:
			continue
		sumtf += count
		words.append( (word,count, flex) )
	return words, sumtf

def load_wn_words(filename, encoding = 'utf8'):
	words = set([])
	for line in codecs.open(filename, 'r', encoding):
		line = line.strip()
		if len(line) > 0:
			words.add(line)
	return words


def load_words_from_matrix(mdir, mname):
	words = []
	sumtf = 0
	matrix = pymatrices.LoadSuperMatrix(mdir,mname, False,False)
	features = matrix.GetRowFeatures()
	rows = len(features)

	for row in xrange(rows):
		feature = features[row]
		word = feature.GetLabel().decode('cp1250')
		flex = ""
		if ':' in word:
			word = word[ word.find(':')+1: ]
			flex = word[ : word.find(':') ]
		TF = feature.GetGlobalTermFrequency()
		sumtf += TF
		words.append( (word,TF, flex ) )
	
	print "Words:",len(words)

	return words,sumtf


def draw_histogram(udict, sumtf, bins = 10):
	vals = numpy.zeros( len(udict) )
	index = 0
	for word, count,flex in udict:
		vals[index] = float( count )
		index += 1

	#vals.sort()
	pylab.hist(vals, bins=bins, normed=False)
	pylab.show()

def select_words(udict, num, bins, wnwords):
	random.seed()
	words = []
	selected_words = set([])
	for lower,upper in bins:
		selected = 0
		while selected < num:
			word,count,flex = udict[random.randint(0, len(udict)-1)]
			if (not word in wnwords) or (word in selected_words):
				continue
			if count >= lower and count < upper:
				selected += 1
				words.append(  (word,count,flex) )
				selected_words.add(word)
	return words

def write_words(name, words, encoding='cp1250'):
	filedict = codecs.open(name+"_dict.txt",'w', encoding)
	fileplain = codecs.open(name+"_plain.txt",'w', encoding)

	for word,count,flex in words:
		line = "%s := %s\n" % (word,count)
		filedict.write( line )
		fileplain.write(word + '\n')

	filedict.close()
	fileplain.close()

if __name__ == '__main__':
	udictfile = '/mnt/data/in/udict/udict_guesser_kipi.txt'
	matrixdir = '/mnt/data/workdir/macierze/supermatrix2/noun/'
	wnwordsfile = '/home/boban/subst_plwordnet_all_robocza.txt'
	matrixname=  'kipi_substAdjP'
	print "Loading udict..."
	dict,tf = load_words_from_matrix(matrixdir, matrixname)
	print "Size of dict: %s\nLoading words from wordnet" % len(dict)
	wnwords = load_wn_words( wnwordsfile )
	bins = [ (4,50), (50, 100), (100,250), (250,500), (500,750), (750, 1000), (1000,10000), (10000, sys.maxint) ]
	print "Loaded %s words from wordnet file\nSelecting words..." % len(wnwords)
	selected = select_words(dict, 10, bins, wnwords)
	print "Writing results..."
	write_words('selected_words', selected)
	print "Done..."




