# -*- coding: utf-8 -*-
import subprocess, codecs, sys, os, scipy.sparse, math

class MiniPar:
	"""
	Reads and parses single document.
	"""
	def __init__(self, minpar_data_path = '/home/boban/src/soft/minipar/data/', \
				minpar_pdemo_path = '/home/boban/src/soft/minipar/pdemo/pdemo',
				words = 'words_bnc_sorted.txt'):
		self.minipar_data = minpar_data_path
		self.minipar_pdemo = minpar_pdemo_path
		self.triples = {} # (w1,rel,w2) -> count
		self.w2_cat = set(["N", "V", "SaidX", "XSaid", "A"])
		self.relations = set(["s",'vrel','obj','obj2','obj1', 'desc', 'gen', 'nn', 'conj', 'amod', 'person', 'city', 'mod', 
							'appo', 'mod-before', 'sc' 'amount-value', 'pnmod', 'inside', 'lex-dep', 'as-arg'])
		
		self.bases = {} # (category,base) -> count
		self.nouns = set([]) # only base forms
		self.all_words = {} # category,orth -> base 
		 
		self.load_words(words)
		
		# data from http://www.natcorp.ox.ac.uk/docs/userManual/design.xml
		# TODO: should check this manually
		self.lexical_units = 97619934
		
		# filled during saving matrix:		
		self.colum_desc = {} # (relation, word) -> column_id
		self.used_nouns = {} # noun -> row_id
		
		#TODO: stoplista
		
	def load_words(self, filename):
		for line in codecs.open(filename, 'r', 'utf8'):
			line = line.strip();
			if len(line) == 0: continue
			tokens = line.split(';')
			base = tokens[0]
			category = tokens[1]
			count = tokens[2]
			orths = set(tokens[3:])
			self.bases[(category, base)] = count
			if category == "N":
				self.nouns.add(base)
			for orth in orths:
				self.all_words[ (category, orth) ] = base 
		print "Loaded %d bases and %d orths (%d nouns)" % (len(self.bases), len(self.all_words), len(self.nouns) )
			
	def extract_triplets(self, filename):
		p1 = subprocess.Popen(["cat", filename], stdout=subprocess.PIPE)
		p2 = subprocess.Popen([self.minipar_pdemo, '-t', '-i', '-p', self.minipar_data], stdin=p1.stdout, stdout=subprocess.PIPE)
		for line in p2.stdout:
			line = line.strip().decode('utf8')
			if len(line) == 0: continue	
			tokens = line.split('\t')
			if len(tokens) != 3:
				print "Wrong pare line:", line
				continue
			self.add_triples(tokens)
			# try to add it in reverse order
			tokens.reverse()
			rel = tokens[1].split(':')
			rel.reverse()
			rel[1] = rel[1] + '-of'
			rel = ':'.join(rel)
			tokens[1] = rel
			self.add_triples(tokens)

	def get_word_base(self, word, cat):
		pair = (cat, word)
		if pair in self.bases:
			return pair
		if not pair in self.all_words:
			return None
		return cat,self.all_words[pair]

	def add_triples(self, triple):
		tokens = triple[1].split(':')
		w1_str = triple[0]
		w1_cat = tokens[0] 
		if w1_cat == 'NN': w1_cat = 'N' # NN also
		rel = tokens[1]
		w2_str = triple[2]
		w2_cat = tokens[2] #[0] # all verbs types
		if w2_cat == "Aux":
			w2_cat = "V" # auxiliary verbs
		elif w2_cat == 'NUM':
			return
		else:
			w2_cat = w2_cat[0] # first letter for verbs 						
		
		if w1_cat != 'N' or not w2_cat in self.w2_cat: return
		 
		if 'subj' in rel or rel in self.relations:
			w1 = self.get_word_base(w1_str, w1_cat)
			w2 = self.get_word_base(w2_str, w2_cat)
			if w1 == None or w2 == None:
				return	
			triple = (w1,rel,w2) 
			if triple in self.triples:
				self.triples[triple] += 1
			else:
				self.triples[triple] = 1
				
	def save_triples(self, filename):
		file = codecs.open(filename, 'w', 'utf8')
		file.write("w1;w1cat;rel;w2;w2cat;count\n")
		for t in self.triples:
			count = self.triples[t]
			cat1,w1 = t[0]
			r = t[1]
			cat2,w2 = t[2]
			line = "%s;%s;%s;%s;%s;%d\n" % (w1, cat1, r, w2, cat2, count)
			file.write(line)
		file.close()  
	
	def load_triples_from_files(self, files):
		self.triples.clear()
		for filename in files:
			self.load_triples(filename)
	
	def load_triples(self, filename):
		first = True
		for line in codecs.open(filename, 'r', 'utf8'):
			line = line.strip();
			if first:
				first = False
				continue
			
			if len(line) == 0: continue
			tokens = line.split(';')
			w1 = (tokens[1],tokens[0])
			rel = tokens[2]
			w2 = (tokens[4],tokens[3])
			count = tokens[5]
			triple = (w1,rel,w2)
			if triple in self.triples:
				self.triples[triple] += count
			else:
				self.triples[triple] = count
				
		print "Loaded %s triples from %s" % (len(self.triples), filename )
		
	def __prapare_dir(self, directory, matrixname):
		outdir = os.path.join(directory, matrixname)
		if os.path.exists(directory):
			if not os.path.isdir(directory):
				raise IOError, "Path %s exists and is not a directory!" % directory
			if os.path.exists(outdir):
				raise IOError, "Directory (%s) for matrix exists!" % outdir
			os.makedirs(outdir)				
		else:
			os.makedirs(outdir)
			if not os.path.exists(directory):
				raise IOError, "Failed to create %s" % outdir
			
		return outdir
	
	def __create_column_desc(self):
		# first gather data about all possible columns
		data = {} # relation -> set(words) 
		for t in self.triples:
			r = t[1]
			if r in data:
				data[r].add(t[2])
			else:
				data[r] = set([t[2]])				
							
		# store it sorted by relation
		last_id = 0
		self.column_words = []
		for relation in sorted(data):
			words = data[relation]
			for w in words:
				self.colum_desc[ (relation, w) ] = last_id
				self.column_words.append( (relation,w) )
				last_id += 1

		print "Found %s columns for %s relations" % (last_id, len(data))
		
	def __create_row_desc(self):
		self.row_nouns = set([])
		for t in self.triples:
			count = self.triples[t]
			cat1,w1 = t[0]
			if cat1 != 'N': continue
			self.row_nouns.add(w1)
		
		self.row_nouns = sorted(self.row_nouns)
		last_id = 0
		for noun in self.row_nouns:
			if not noun in self.used_nouns:
				self.used_nouns[noun] = last_id
				last_id += 1		
		
		print "Found %d row nouns" % last_id 			
				
	def __write_crs(self, mdir):
		print "Preparing row/lil_matrix"
		self.crs = scipy.sparse.lil_matrix( (len(self.used_nouns), len(self.colum_desc) ) )
		triples_processed = 0
		for t in self.triples:
			#self.__create_row_for_noun(noun)
			cat1,w1 = t[0]
			if cat1 != 'N': continue
			r = t[1]
			count = int(self.triples[t])
			
			column_no = self.colum_desc[(r, t[2])]
			row_no = self.used_nouns[w1]	
			
			self.crs[row_no, column_no] = self.crs[row_no, column_no] + count
			if triples_processed % 10000 == 0:
				print "Triples processed %s / %s" %(triples_processed, len(self.triples))
			triples_processed+=1
		
		shape = self.crs.get_shape()
		self.nnz = self.crs.getnnz()
		
		filename = os.path.join(mdir, 'matrix_crs.mat')
		print "Writing crs part to %s, matrix has %d nnz (%f%%) " % (filename, self.crs.getnnz(), \
																	 (self.crs.getnnz() / float(shape[0]*shape[1]))*100 )
		file = open(filename, 'w')
		file.write('%d %d %d\n' % (shape[0], shape[1], self.crs.getnnz() ) )
		for row_no in xrange(shape[0]):
			if len(self.crs.rows[row_no]) == 0:
				file.write('\n')
			else:
				ind = self.crs.rows[row_no]
				data = self.crs.data[row_no]
				line = ''
				for i in xrange(len(ind)):
					tuple = '%d %d ' % (ind[i], data[i])
					line += tuple
				file.write(line.strip())
				file.write('\n')
		file.close()
	
	def __write_csc(self, mdir):
		print "Preparing column/lil_matrix"
		# inverted matrix, rows are columns 
		self.ccs = scipy.sparse.lil_matrix( (len(self.colum_desc), len(self.used_nouns) ) )
		triples_processed = 0
		for t in self.triples:
			#self.__create_row_for_noun(noun)
			cat1,w1 = t[0]
			if cat1 != 'N': continue
			r = t[1]
			count = int(self.triples[t])
			
			column_no = self.colum_desc[(r, t[2])]
			row_no = self.used_nouns[w1]	
			
			self.ccs[column_no, row_no] = self.ccs[column_no, row_no] + count
			if triples_processed % 10000 == 0:
				print "Triples processed %s / %s" %(triples_processed, len(self.triples))
			triples_processed+=1
			
		filename = os.path.join(mdir, 'matrix_ccs.mat')
		file = open(filename, 'w')
		shape = self.ccs.get_shape()
		print "Writing crs part to %s, matrix has %d nnz (%f%%) " % (filename, self.ccs.getnnz(), \
																	 (self.ccs.getnnz() / float(shape[0]*shape[1]))*100 )
		
		file.write('%d %d %d\n' % (shape[0], shape[1], self.ccs.getnnz() ) )
		for row_no in xrange(shape[0]):
			if len(self.ccs.rows[row_no]) == 0:
				file.write('\n')
			else:
				ind = self.ccs.rows[row_no]
				data = self.ccs.data[row_no]
				line = ''
				for i in xrange(len(ind)):
					tuple = '%d %d ' % (ind[i], data[i])
					line += tuple
				file.write(line.strip())
				file.write('\n')
		file.close()
	
	def __count_entropy(self, row, sum):
		data = row.data[0]
		entropy = 0.0
		for elem in data:
			p = elem / float(sum)
			entropy += p * math.log(p)
			
		return -entropy
	
	def __write_row_features(self, mdir, mname):
		filename_features = os.path.join(mdir, "row_features.csv")
		filename_rlabel = os.path.join(mdir, mname + ".rlabel")
		
		print "Writing row features to %s and rlabel to %s " % (filename_features, filename_rlabel)
		
		file_features = open(filename_features, 'w')
		file_rlabel = open(filename_rlabel, 'w')
		
		file_features.write("label;nz;tf;TF;entropy\n")
		rows = self.crs.get_shape()[0]
		for row_no in xrange(rows):
			row = self.crs.getrowview(row_no)
			nz = row.getnnz()
			tf = int(row.sum())
			label = self.row_nouns[row_no]
			
			TF = self.bases[ ('N', label) ]
			entropy = self.__count_entropy(row, tf)
			
			line = "%s;%s;%s;%s;%s\n" % (label, nz, tf, TF, entropy)
			file_features.write(line)
			file_rlabel.write(label + '\n')
		file_features.close()
		file_rlabel.close()
		
	def __write_column_features(self,mdir, mname):
		filename_features = os.path.join(mdir, "column_features.csv")
		filename_clabel = os.path.join(mdir, mname + ".clabel")
		print "Writing column features to %s and clabel to %s " % (filename_features, filename_clabel)
		
		file_features = open(filename_features, 'w')
		file_clabel = open(filename_clabel, 'w')
		file_features.write("label;nz;tf;TF;entropy\n")
		columns = self.ccs.get_shape()[0]
		for column_no in xrange(columns):
			# getrowview, because this is transposed matrix
			column = self.ccs.getrowview(column_no)
			nz = column.getnnz()
			tf = int(column.sum())
			relation,word = self.column_words[column_no]
			label = relation + ":" +  word[1] + "_" + word[0]
			
			TF = self.bases[ word ]
			entropy = self.__count_entropy(column, tf)
			
			line = "%s;%s;%s;%s;%s\n" % (label, nz, tf, TF, entropy)
			file_features.write(line)
			file_clabel.write(label+ '\n')
			
		file_features.close()
		file_clabel.close()
		
	def __write_info(self, mdir):
		filename = os.path.join(mdir, 'info.txt')
		print "Writing info to", filename
		file = open(filename, 'w')
		file.write("The British National Corpus: World Edition\n")
		file.write(str(self.lexical_units) + "\n")
		file.close()
				 
	def save_supermatrix(self, directory, matrixname):
		# prepare dir
		mdir = self.__prapare_dir(directory, matrixname)
		#create row and column indices
		self.__create_column_desc()
		self.__create_row_desc()		
		# prepare and write crs part
		self.__write_crs(mdir)
		# prepare and write ccs part
		self.__write_csc(mdir)
		# prepare and write row features and .rlabel
		self.__write_row_features(mdir, matrixname)
		# prepare and write column features and .clabel
		self.__write_column_features(mdir, matrixname)
		# write info
		self.__write_info(mdir)
		# write boundary.txt
		# write udict
		# write triples 
		pass
	

def parse_single(minipar, dirname, names):
	for name in names:	
		filename = os.path.join(dirname, name)
		if os.path.isfile(filename):
			print "Processing",filename
			minipar.extract_triplets(filename)

if __name__ == '__main__':
	minipar = MiniPar()
	#print "Processing dir", sys.argv[1]
	#os.path.walk(sys.argv[1], parse_single, minipar)
	#print "Writing triples to", sys.argv[2]
	#minipar.save_triples(sys.argv[2])
	
	minipar.load_triples(sys.argv[1])
	print "Writing supermatrix %s to %s" % (sys.argv[3], sys.argv[2])
	minipar.save_supermatrix(sys.argv[2], sys.argv[3])
	
	
