
def load_udict(filename, encoding = "cp1250"):
	"""
	Reads unigram dictionary in format:
	flex_class:word := TF
	"""
	print "Loading udict from file %s, assuming encoding %s." \
			% (filename, encoding)
	import codecs
	# lemma -> [ flex,tf ] (if tf > 0 )
	udict = {}
	for line in codecs.open(filename, "r", encoding):
		line = line.strip()
		first_split = line.split(":=")

		if len(first_split) != 2: continue
		flex_word = first_split[0].strip()
		
		second_split = flex_word.split(":")

		if len(second_split) != 2: continue

		word = second_split[1].strip()
		words_flex = second_split[0].strip()
		tf = int( first_split[1].strip() )

		if tf == 0:
			print "TF = 0 for word '%s'" % word

		if word in udict:
			udict[word].add(words_flex)
		else:
			udict[word] = set([words_flex])

	print "Loaded %s words from udict" % len(udict)
	
	return udict

def load_words(filename, encoding = "cp1250"):
	"""
	Loads words from file. File should contain one word in every line.
	"""
	import codecs
	print "Loading words from file %s, assumming encoding %s" % ( filename, encoding )
	words = []
	for line in codecs.open(filename, "r", encoding):
		line = line.strip()
		if len(line) > 0:
			words.append( line )
	
	print "Loaded %s words..." % len(words)

	return words

def filter(words, udict, flex):
	"""
	Return part of words, that have given flex. Chcecking of flex class 
	is preformed on the basis of udict.
	"""
	out = []
	for word in words:
		if word in udict:
			flexes = udict[word]
			if flex in flexes:
				out.append(word)

	return sorted(out)

def save_to_file(words, filename, encoding = "cp1250"):
	"""
	Saves words to file
	"""
	import codecs
	file = codecs.open(filename, "w", encoding)
	for w in words:
		file.write(w + "\n")

	file.close()

		
if __name__ == "__main__":
	print "Filtering nouns!"
	udict_file = "/mnt/pwr/in/udicts_wrong_premorphing/udict_kipi_rzepa_wiki.txt"
	words_file = "/mnt/pwr/in/slowa/wiersze/rzeczowniki.txt"
	words = load_words(words_file)
	udict = load_udict(udict_file)

	deprs = filter(words, udict, "depr")
	print "Found %s deprs" % len(deprs)
	
	gers = filter(words, udict, "ger")
	print "Found %s gers" % len(gers)

	substs = filter(words, udict, "subst")
	print "Found %s substs" % len(substs)

	save_to_file(deprs, "all_depr.txt")
	save_to_file(gers, "all_gers.txt")
	save_to_file(substs, "all_subst.txt")

	all = set(words)
	all.difference_update( set(substs) )
	all.difference_update( set(deprs) )
	all.difference_update( set(gers) )
	
	print "Not found %s words" % len(all)
	if len(all):
		save_to_file( all, "not_found.txt")

