#!/usr/bin/python
# -*- coding: utf-8 -*-

from HWBST import WBSTQuestion, HWBSTGenerator
from scipy import sparse
import random


class EWBSTGenerator(HWBSTGenerator):
	"""
	"""
	
	def __init__(self, db, wn, candidate_words, detractor_words, no_of_detractors = 3, close_path = 4, 
			add_hypernyms = True, extend_small = True, exclude_neighbours = True, log = False):
		"""
		
		Params:
		db - pyplwn.Database object
		wn - initalized pyplwn.WNGraph object - for counting paths between words
		candidate_words - words only for questions
		detractor_words - words for detractors only
		no_of_detractors - how many detractors are generated for each QA pair
		close_path - longest path between question (or answer) and detractor
		extend_small - if True then when there will be not enough detractor within (path) range,  
			then detractors will be extended with:
			1. connected words, even if path length between detractor and question (or answer) 
				will be grater then close_path
			2. not connected words

		"""
		HWBSTGenerator.__init__(self, db, candidate_words, detractor_words, no_of_detractors, add_hypernyms)
		self.close_path = close_path
		self.wn = wn
		self.extend_small = extend_small
		self.log = log
		self.exclude_neighbours = exclude_neighbours
		
	def __generate_matrix(self):
		print "Generating length matrix..."
		all_words = []
		all_words.extend(self.detractor_words)
		all_words.extend(self.candidate_words)
		all_words_set = set(all_words)
		all_words = []
		all_words.extend(all_words_set)
		del all_words_set
		self.path_matrix = WNPathMatrix(self.wn, all_words)	
		print "Generating length matrix... Done."

	def generate(self):
		"""
		"""
		self.__generate_matrix();
		self.extended_connected = 0
		self.extended_no_connected = 0
		self.close_only = 0

		questions = []
		for w1_ind in xrange(len(self.candidate_words)):
			w1 = self.candidate_words[w1_ind]
			# get close words - possible answers
			close_words = self.get_all_words_with_words(w1)
			print "Processing word " + str(w1_ind) + " of " + str(len(self.candidate_words))

			for w2 in close_words:
				if not w2 in self.detractor_words:
					continue
				if not self.can_add_question(w1,w2):
					continue
				
				# extend close words with those that are close to answer
				all_close_words = close_words.union( self.get_all_words_with_words(w2) )
				# generate detractors
				detractors = self.generate_detractors( w1,w2, all_close_words )
				if len(detractors) == self.no_of_detractors:
					if w1 == 'stolec' or w1 == 'kupa':
						print WBSTQuestion(w1, w2, detractors).to_csv_string()
					questions.append( WBSTQuestion(w1, w2, detractors) )
				elif self.log: 
					print "not enough detractors [%s] for: (%s,%s), generated %s so far [ %s/ %s ]" \
							% (len(detractors), w1,w2, len(questions), w1_ind, len(self.candidate_words) )

			if w1_ind % 100 == 0:
				print "Generating ewbst test for word %s (%s / %s). Generated %s questions so far..." % \
						(w1, w1_ind, len(self.candidate_words), len(questions))

		print "Generated %s questions for %s candidate words" % (len(questions), len(self.candidate_words) )
		print "Extension statistics...\n  Generated without extending: %s\n  Generated with extension of connected words: %s\n  Generated with extension of non-connected words: %s" % ( self.close_only, self.extended_connected, self.extended_no_connected)
		return questions


	def generate_close_detractors(self, w1, w2, close_words):
		"""
		Generate detractors for q/a pair (w1,w2) that are in range (self.close_words), 
		but are not in close_words.
		"""
		detractors = []
		# all possible detractors --- words that are connected with w1 or w2 and are in self.detractor_words
		possible_detractors = []
		possible_detractors.extend( ( set(self.path_matrix.get_nz_words( w1 )).\
				union(self.path_matrix.get_nz_words(w2) ) ).intersection(self.detractor_words) )
		random.shuffle(possible_detractors)

		# search for proper detractors
		while len(detractors) < self.no_of_detractors and len(possible_detractors) > 0:
			w = possible_detractors.pop()
			# too close words are not good for detractors
			if w == w1 or w == w2 or w in close_words:
				continue
			# we need only words from detractor_words and wordnet lemmas
			if not w in self.detractor_words or not w in self.lemmas:
				continue
			path_len_w1 = self.path_matrix.path(w1,w)
			path_len_w2 = self.path_matrix.path(w2,w)
			
			# if we are excluding close neighbours - filter it now
			if self.exclude_neighbours and (path_len_w1 == 1 or path_len_w2 == 1):
				continue

			# check if selected word is close enough to be included in detractors

			if path_len_w1<=0:
				path_len_w1=9001
			if path_len_w2<=0:
				path_len_w2=9001
			path_len = min(path_len_w1, path_len_w2)
			if path_len < self.close_path:
				detractors.append(w)
				#print "APPENDING " + w + " to " + w1 + " detractors with path " + str(path_len) + "<" + str(self.close_path)
			else:
				#rint "NOT APPENDING " + w + " to " + w1 + " detractors with path " + str(path_len) + ">=" + str(self.close_path)
				pass

		return detractors

	def extend_with_connected(self, w1, w2, close_words, detractors):
		"""
		Extends detractors with words that are connected to w1 or w2, but not necessary in close_path range.
		"""
		possible_detractors = []
		# possible detractors: union of set of words connected to w1 or w2, 
		# but only those words that are in self.detractor_words, 
		#without words already in detractors 
		possible_detractors.extend( (( set(self.path_matrix.get_nz_words( w1 )).\
				union(self.path_matrix.get_nz_words(w2) ) ).intersection(self.detractor_words)).difference(set(detractors)) )
		random.shuffle(possible_detractors)

		# search for proper detractors
		while len(detractors) < self.no_of_detractors and len(possible_detractors) > 0:
			w = possible_detractors.pop()
			# too close words are not good for detractors
			if w == w1 or w == w2 or w in close_words:
				continue
			# we need only words from detractor_words and wordnet lemmas
			if not w in self.detractor_words or not w in self.lemmas:
				continue

			path_len_w1 = self.path_matrix.path(w1,w)
			path_len_w2 = self.path_matrix.path(w2,w)
			
			# if we are excluding close neighbours - filter it now
			if self.exclude_neighbours and (path_len_w1 == 1 or path_len_w2 == 1):
				continue

			detractors.append(w)
	
	def extend_with_no_connected(self, w1, w2, close_words, detractors):
		possible_detractors = []
		# possible detractors: (all detractors \ (detractors + close_words ) ) n all_lemmas
		possible_detractors.extend( (set(self.detractor_words).difference(set(detractors).union(set(close_words)) ) ).intersection(self.lemmas) )
		random.shuffle(possible_detractors)

		# search for proper detractors
		while len(detractors) < self.no_of_detractors and len(possible_detractors) > 0:
			w = possible_detractors.pop()
			if w == w1 or w == w2 :
				continue

			# if we are excluding close neighbours - filter it now
			if self.exclude_neighbours:
				path_len_w1 = self.path_matrix.path(w1,w)
				path_len_w2 = self.path_matrix.path(w2,w)
				if path_len_w1 == 1 or path_len_w2 == 1:
					continue

			detractors.append(w)


	def generate_detractors(self, w1, w2, close_words):
		"""
		Generates detractors for qa pairs (w1 and w2 params), 
		close_words is a collection of words that are to w1 or w2 close to be detractors
		"""
		detractors = self.generate_close_detractors(w1,w2,close_words)

		if len(detractors) < self.no_of_detractors and self.extend_small:
			self.extend_with_connected( w1, w2, close_words, detractors)
			# if still not enough detractors, but we have at least one (to ensure difficulty of the test)
			if len(detractors) > 0 and len(detractors) < self.no_of_detractors:
				self.extend_with_no_connected(w1, w2,close_words, detractors)
				if len(detractors) == self.no_of_detractors:
					self.extended_no_connected += 1
			elif len(detractors) == self.no_of_detractors:
				self.extended_connected += 1
		elif len(detractors) == self.no_of_detractors:
			self.close_only += 1

		return detractors
		
class WNPathMatrix(object):
	def __init__(self, wngraph, words):
		"""
		Params:
		wngraph - pyplwn.WNGraph object
		words - a collection of words for which the matrix will be built
		"""
		self._init_words(words)
		self._init_matrix(wngraph)

	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):
		"""
		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)

		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]

			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 < 11:
				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

	
	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
		"""
		print "NZ WORDS CALL"
		if not self.words_to_row.has_key(word):
			return []

		row = self._matrix.rows[ self.words_to_row[ word ] ]
		print "NZ ROW: " + str(row)

		words = []
		for v in row:
			words.append( self.words[v] )
		print "NZ WORDS RET"
		return words

