#!/usr/bin/python
# -*- coding: utf-8 -*-
import pyplwn, random, wordsutils, os
from scipy import sparse

class RowGenerator(object):
	def __init__(self, wngraph, words, row):
		"""
		Params:
		wngraph - pyplwn.WNGraph object
		words - a collection of words for which matrix will be built
		"""
		self.row = row
		self.__init_words(words)
		self.__init_matrix(wngraph,11)

	def __init_words(self, words):
		"""
		Prepares two collections:
		self.word_to_row - a mapping between words and number of row
		self.words - a list of words
		"""
		words = set(words)
		self.words_to_row = {}
		self.words = []
		id = 0
		for w in words:
			self.words.append(w)
			self.words_to_row[w] = id
			id += 1
	
	def __init_matrix(self, graph, len_limit = 11):
		"""
		Prepares self.__matrix, fill cells of matrix with path lengths between words.
		Because |path(a,b)| == |path(b,a)| only values above diagonal will be stored.
		Code based on pyplwn.utils.extract_pairs_for_relation_min_max_len by Czuk

		"""
		self.__matrix = sparse.lil_matrix( (len(self.words), len(self.words) ) )
		import Queue, networkx
		
		lemmas4synset_id = graph.db.Synsets.get_units_lemma_for_synsets_id()

		# Słownik (lemat, lemat):minimalna odległość
		pairs_length = {}
		pairs_queue = Queue.Queue()
		pairs_index = 0

		all_words = set([self.words[self.row]])

		print len(lemmas4synset_id)*len(lemmas4synset_id)
		# Dodajemy synonimy
		for sid in lemmas4synset_id:
			for unit1 in lemmas4synset_id[sid]:
				for unit2 in lemmas4synset_id[sid]:
					pairs_length[(unit1, unit2)] = 0

		# Dodajemy wszystkie pary o długości 1
		for edge in graph.graph.edges_iter():			
			if lemmas4synset_id.has_key(edge[0]) and lemmas4synset_id.has_key(edge[1]):# and graph.db.Synsets.get_synset_pos(edge[0])==2:
				lemmas1 = lemmas4synset_id[edge[0] ]
				lemmas2 = lemmas4synset_id[edge[1] ]
				lemmas = lemmas1.union(lemmas2)

				if len( lemmas.intersection( all_words) ) > 0:
					pairs_queue.put((edge,1))

		count = 1;
		while not pairs_queue.empty():
			tuple = pairs_queue.get()

			# Dodaj pary lematów z odległościami
			added = False
			edge = tuple[0]
			length = tuple[1]

			if length <= len_limit:
				for u1 in lemmas4synset_id[edge[0]]:
	#				if not self.words_to_row.has_key(u1):
	#					continue
					for u2 in lemmas4synset_id[edge[1]]:
	#					if not self.words_to_row.has_key(u2):
	#						continue
						if not pairs_length.has_key((u1,u2)) and u1!=u2:
							pairs_length[(u1,u2)] = length
							added = True
			
			if added and length < len_limit:
				for edge_child in graph.graph.edges(edge[1]):
					if lemmas4synset_id.has_key(edge_child[1]):
						pairs_queue.put(((edge[0],edge_child[1]),length+1))

			if count % 5000 == 0:
				print pairs_queue.qsize()
			count += 1
		

		for edge in pairs_length:
			length = pairs_length[edge]
			w1 = edge[0]
			w2 = edge[1]
			if not self.words_to_row.has_key(w1) or not self.words_to_row.has_key(w2):
				continue
			w1_ind = self.words_to_row[w1]
			w2_ind = self.words_to_row[w2]
			self.__matrix[w1_ind, w2_ind] = length
			if length<2 and length>0:
				print "ADDING " + w1 + " " + w2 + " with len " + str(length)

	
	def path(self, word1, word2):
		"""
		Return length of path between to words. If there was no path then returns -1.
		If there is no data for word1 or word2 then returns -2.
		"""
		if not self.words_to_row.has_key(word1) or not self.words_to_row.has_key(word2):
			return -1
		if word1 == word2:
			return 0

		word1 = self.words_to_row[word1]
		word2 = self.words_to_row[word2]
		if word1 < word2:
			path_len = self.__matrix[ word2 , word1  ]
		else:
			path_len = self.__matrix[ word1 , word2  ]

#		return path_len 
		return self.__matrix[word1, word2]

	def get_nz_words(self, word):
		"""
		Returns list of words with path > 0 for given word
		"""
		if not self.words_to_row.has_key(word):
			return []

		row = self.__matrix.rows[ self.words_to_row[ word ] ]

		words = []
		for v in row:
			words.append( self.words[v] )

		return words

	def save_row(self, filename):
		print "Saving row for " + self.words[self.row]
		ofile = open(filename, "w")
		for i in xrange(0,len(self.__matrix.rows)):
			if (self.__matrix[self.row, i] > 0):
				ofile.write(str(i) + " " + str(self.__matrix[self.row, i]) + " ")


def main(wordsfile, udictfile, outdir, row_index):
	import pickle
	db = pyplwn.database.Database(user = "wordnet", passwd = "mylocaldbpasswd", host = "127.0.0.1", db = "wordnet", port = 3306, encoding = "utf8")
	db.set_pos_filter([pyplwn.database.DBEnums.get_pos('subst')])

	if row_index == 0:
		all_words = wordsutils.load_words(wordsfile)
		print "Before filtering:", len(all_words)
		all_words = wordsutils.wordnet_filter( all_words, db )
		print "After filtering:", len(all_words)
		wn = pyplwn.wngraph.WNGraph( [db.Relations.get_relation_by_name("hiperonimia"), db.Relations.get_relation_by_name("hiponimia")], db, False )
		matrixrow = RowGenerator(wn, all_words, 0)
		matrixrow.save_row(os.path.join(outdir, "0"))
		ofile = open(os.path.join(outdir, "wngraph"), "wb")
		pickle.dump(wn.graph,ofile,pickle.HIGHEST_PROTOCOL)
		ofile = open(os.path.join(outdir, "words"), "wb")
		pickle.dump(all_words,ofile,pickle.HIGHEST_PROTOCOL)
		ofile = open(os.path.join(outdir, "words2index"), "wb")
		pickle.dump(matrixrow.words_to_row,ofile,pickle.HIGHEST_PROTOCOL)
		print "I HAS ME " + str(len(all_words)) + " WORDS."
	else:
		ifile = open(os.path.join(outdir, "words"), "rb")
		all_words = pickle.load(ifile)
		wn = pyplwn.wngraph.WNGraph( [db.Relations.get_relation_by_name("hiperonimia"), db.Relations.get_relation_by_name("hiponimia")], db, False )
		ifile = open(os.path.join(outdir, "wngraph"), "rb")
		wn.graph = pickle.load(ifile)
		matrixrow = RowGenerator(wn, all_words,row_index)
		matrixrow.save_row(os.path.join(outdir, str(row_index)))

if __name__ == "__main__":
	import sys, argparse
	parser = argparse.ArgumentParser(description='Row matrix generator')
	#positional arguments
	parser.add_argument('words_file',  action='store', 
		help='Path to the file with words ')
	parser.add_argument('dict_file',  action='store', 
		help='Path to the file with dictionary')
	parser.add_argument('out_dir',  action='store', 
		help='Path to output directory')
	parser.add_argument('row_index',  action='store', 
		help='Number of row')
	
	
	args = parser.parse_args()
	
	main(args.words_file, args.dict_file, args.row_index)
