# -*- coding: utf-8 -*-
"""
Converts dependency triples into supermatrix.

Author: Bartosz Broda, bartosz.broda@gmail.com
"""

import os,sys,codecs, math

from scipy.sparse import lil_matrix

import matrix_writer

def gather_data(files, dictionary):
	"""
	Gathers data (indices) from the files, but uses only words that are in the dictionary.
	"""
	rows_term_indices = set([]) # indices in dictionary
	column_rel_term = {} # relation -> set([termids])
	for i,filename in enumerate(files):
		print "Indexing %d / %d (%s)" % (i, len(files), filename)
		with codecs.open(filename, 'r', 'utf8') as f:
			for line in f:
				line = line.strip()
				if len(line):
					l,rel,r, c = [x.strip() for x in line.split(';')]
					if l in dictionary and r in dictionary:
						rows_term_indices.add( dictionary[l] )
						if rel in column_rel_term:
							column_rel_term[rel].add( dictionary[r] )
						else:
							column_rel_term[rel] = set([ dictionary[r] ])
	return rows_term_indices, column_rel_term

def create_matrix(files, dictionary):
	"""
	Creates a matrix from a given set of data files which contain tuples of 
	left; relation; right; count

	Filters left and right using a dictionary.
	"""
	rows_term_indices, column_rel_term = gather_data(files, dictionary)

	# index
	term2row_no = dict( [ (x[1], x[0]) for x in enumerate(rows_term_indices)] )
	termrel2col_no = {}
	index = 0
	for rel in column_rel_term:
		for termid in column_rel_term[rel]:
			termrel2col_no[ (rel, termid) ] = index
			index += 1

	rows = len(term2row_no)
	cols = len(termrel2col_no)
	matrix = lil_matrix( (rows, cols) )
	for i,filename in enumerate(files):
		print "Processing %d / %d (%s)" % (i, len(files), filename)
		with codecs.open(filename, 'r', 'utf8') as f:
			for line in f:
				line = line.strip()
				if len(line):
					l,rel,r, c = [x.strip() for x in line.split(';')]
					count = int(c)
					if l in dictionary and r in dictionary:
						row_index = term2row_no[ dictionary[l] ]
						column_index = termrel2col_no[ (rel, dictionary[r]) ]
						matrix[row_index, column_index] += count

	# row labels
	rev_dict = dict( [(dictionary[x],x) for x in dictionary] )
	row_no2termid = dict( [ (term2row_no[x], x) for x in term2row_no])
	matrix.row_labels = []
	for row_no in xrange(rows):
		matrix.row_labels.append( rev_dict[ row_no2termid[row_no] ] )
	del row_no2termid
	# column labels
	matrix.column_labels = []
	matrix.boundaries = {}
	col_no2termrel = dict( [ (termrel2col_no[x], x) for x in termrel2col_no])
	last_rel = col_no2termrel[0][0]
	for col_no in xrange(cols):
		(rel, termid) = col_no2termrel[col_no]
		term = rev_dict[termid]
		matrix.column_labels.append( '%s:%s' % (rel, term) )
		if rel != last_rel:
			matrix.boundaries[col_no+1] = last_rel
	matrix.boundaries[col_no+1] = last_rel
	return matrix

def count_features(matrix,labels, tfdictionary):
	"""
	Count matrix features needed for writing the matrix in SuperMatrix format.

	Arguments:
	- matrix a lil-like  matrix
	- labels - a dictionary of labels for a given matrix
	- tfdictionary - an unigram dictionary (word -> frequency)
	"""
	def count_entropy(data, tf):
		entropy = 0.0
		for p in [x / float(tf) for x in data]:
			entropy += p * math.log(p)
		if entropy != 0:
			return -entropy
		return 0

	features = [] # (label,nz,tf,TF,entropy)
	rows = matrix.shape[0]
	for row_no in xrange(rows):
		data = matrix.data[row_no]
		label = labels[row_no]
		nz = len(data)
		tf = sum(data)
		if ':' in label:
			TF = tfdictionary[ label[ label.index(':')+1:] ]
		else:
			TF = tfdictionary[ label ] 
		if TF == 0:
			sys.stderr.write("Warning: no TF for %s, using tf = %d\n" % (label, tf) )
			TF = tf
		entropy = count_entropy(data, tf)
		features.append( (label, nz, tf, TF, entropy) )

	return features


def main():
	import argparse
	parser = argparse.ArgumentParser(description='Converts directory with dependency triples into a SuperMatrix')
	parser.add_argument('-i', '--indir', dest = 'indir', required=True, 
			help = 'Input directory with dependency tripeles.')
	parser.add_argument('-o', '--outdir', dest = 'outdir', default = './', 
			help = 'Output directory for writing SuperMatrix.')
	parser.add_argument('-p', '--prefix', dest = 'prefix', default = 'data', 
			help = 'Prefix for files in input directory, default: data.')
	parser.add_argument('-u', '--udict', dest = 'udict', required=True, 
			help = 'File with unigram dictionary.')
	parser.add_argument('-n', '--name', dest = 'name', required=True, 
			help = 'Matrix name.')
	parser.add_argument('-c', '--comment', dest = 'comment', default = 'no description given' ,
			help = 'Description of a matrix.')
	
	args = parser.parse_args()
	
	with codecs.open(args.udict, 'r', 'utf8') as f:
		lines = f.readlines()
		tokens = int(lines[0])
		tfdictionary = dict( [(y[0], int(y[1])) for y in [ x.strip().split(';') for x in lines[1:]] ] ) 
	dictionary = dict( [ (x[1], x[0]) for x in enumerate(tfdictionary)] )

	allfiles = []
	for x, ddd in sorted([(r,f) for (r,d,f) in os.walk(args.indir) if len(f)] ): 
		allfiles.extend( [os.path.join(x, w) for w in ddd if w.startswith(args.prefix)] )

	allfiles = sorted(allfiles)

	matrix = create_matrix(allfiles, dictionary)

	matrix.row_features = count_features(matrix, matrix.row_labels, tfdictionary)
	matrix.column_features = count_features(matrix.transpose(), matrix.column_labels, tfdictionary)
	matrix.name = args.name
	matrix.comment = args.comment
	matrix.udict = tfdictionary
	matrix.tokens = tokens

	matrix_writer.write(matrix, args.outdir)

	return matrix

if __name__ == '__main__':
	main()
