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

from HWBST import WBSTQuestion
from EWBST import EWBSTGenerator, WNPathMatrix
from scipy import sparse
from wbst.generate_distmatrix_row import RowGenerator
import os, pickle

class LazyEWBSTGenerator(EWBSTGenerator):
	"""
	"""
	
	def __init__(self, db, wn, wordsfile, rowsdir, candidate_words, detractor_words, no_of_detractors = 3, close_path = 4, 
			add_hypernyms = True, extend_small = True, exclude_neighbours = True, log = True):
		"""
		
		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

		"""
		EWBSTGenerator.__init__(self, db, wn, candidate_words, detractor_words, no_of_detractors, close_path,
                                        add_hypernyms, extend_small, exclude_neighbours, log)
		self.rowsdir = rowsdir
		self.wordsfile = wordsfile
		

	def __generate_matrix(self):
		
		print "Not generating length matrix at all."
		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
		if not os.path.isfile(self.wordsfile):
			matrixrow = RowGenerator(self.wn, self.candidate_words, 0)
			matrixrow.save_row(os.path.join(self.rowsdir, "0"))
			ofile = open(os.path.join(self.rowsdir, "wngraph"), "wb")
			pickle.dump(self.wn.graph,ofile,pickle.HIGHEST_PROTOCOL)
			ofile.close()
			ofile = open(os.path.join(self.rowsdir, "words"), "wb")
			pickle.dump(self.candidate_words,ofile,pickle.HIGHEST_PROTOCOL)
			ofile.close()
			o1file = open(os.path.join(self.rowsdir, "words2index"), "wb")
			pickle.dump(matrixrow.words_to_row,o1file,pickle.HIGHEST_PROTOCOL)
			o1file.close()
			print "I HAS ME " + str(len(self.candidate_words)) + " WORDS."
			
		self.path_matrix = WNPathMatrixLazy(self.wordsfile, self.rowsdir, all_words)

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

		questions = []
		
		ifile = open(os.path.join(self.rowsdir, "words"), "rb")
		self.words = pickle.load(ifile)
		words2indfile = open(os.path.join(self.rowsdir,"words2index"),"rb")
		self.words_to_row = pickle.load(words2indfile)
		ifile = open(os.path.join(self.rowsdir, "wngraph"), "rb")
		self.wn.graph = pickle.load(ifile)
		
		for w1_ind in xrange(len(self.candidate_words)):
			self.path_matrix.clearMatrix()
			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) )
				if not os.path.isfile(os.path.join(self.rowsdir, str(self.words_to_row[w1]))):
					matrixrow = RowGenerator(self.wn, self.words ,self.words_to_row[w1])
					matrixrow.save_row(os.path.join(self.rowsdir, str(self.words_to_row[w1])))
				if not os.path.isfile(os.path.join(self.rowsdir, str(self.words_to_row[w2]))):
					matrixrow = RowGenerator(self.wn, self.words ,self.words_to_row[w2])
					matrixrow.save_row(os.path.join(self.rowsdir, str(self.words_to_row[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
		

class lazySparseMatrix(sparse.lil_matrix):
	
	def __init__(self, matrixdir, arg1, shape=None, dtype=None, copy=False):
		self.matrixdir = matrixdir
		sparse.lil_matrix.__init__(self,arg1, shape, dtype, copy)
		self.offset = 100

	def loadMatrixRow(self, rowNumber):
		rowfile = open(os.path.join(self.matrixdir, str(rowNumber)),"r")
		line = rowfile.readline()
		line = line.strip()
		tokens = line.split(" ")
		if len(tokens) % 2 != 0 and line.strip() != "":
			print >> sys.stderr, "len(tokens) % 2 != 0, line:",line
			return
		if line.strip() == "":
			return
		for tok_no in xrange(0, len(tokens), 2):
			id = int(tokens[tok_no])
			val = float(tokens[tok_no+1])
			self[(rowNumber, id)] = val
#			if val >= 0:
#				self[(rowNumber, id)] = val - self.offset
		rowfile.close()
	
	def __getitem__(self, index):
		if len(self.rows[index[0]]) == 0:
			self.loadMatrixRow(index[0])
		return sparse.lil_matrix.__getitem__(self,index)

class WNPathMatrixLazy(WNPathMatrix):

	def __init__(self, wordfile, rowsdir, wordlist):
		ifile = open(wordfile,"rb")
		self.matrixdir = rowsdir
		self.words = pickle.load(ifile)
		words2indfile = open(os.path.join(rowsdir,"words2index"),"rb")
		self.words_to_row = pickle.load(words2indfile)
		self._matrix = lazySparseMatrix(self.matrixdir, (2*len(self.words),2*len(self.words)))
		print "WORDS LEN: " + str(len(self.words))
		
	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]
		#print "INDEX 0 IS " + str(word1)
		return self._matrix[word1, word2]

	def clearMatrix(self):
		del(self._matrix)
		self._matrix = lazySparseMatrix(self.matrixdir, (2*len(self.words), 2*len(self.words)))

	def _init_matrix(self, graph, words):
		pass

	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 []
			
		foo = self._matrix[(self.words_to_row[ word ],0)]
		row = self._matrix.rows[ self.words_to_row[ word ] ]

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

	def tehtype(self):
		print "TYPECHECK"
		print type(self)
		print type(self._matrix)
