import codecs, os, pymatrices, sys

def extract_row_words_from_matrix(matrix, filename, out_encoding = 'utf8', in_encoding = 'utf8'):
	"""
	Extract all row words from matrix and saves them to file. Does no processig on words (i.e. does not remove flex prefix)
	
	Params:
	matrix - SuperMatrix object
	filename - output filename
	out_encoding - output file encoding
	in_encoding - encoding of matrix features
	"""
	rows = matrix.NumberOfRows()
	rowfeatures = matrix.GetRowFeatures()
	file = codecs.open(filename, "w", out_encoding)

	for row_no1 in xrange(rows):
		word1 = rowfeatures[row_no1].GetLabel()
		line = "%s\n" % (word1)
		file.write(line.decode(in_encoding))
	file.close()
		
def extract_stats_for_words_from_matrix(matrix, filename, words, out_encoding = 'utf8', matrix_encoding = 'utf8'):
	"""
	Extract all row words from matrix and saves them to file. Does no processig on words (i.e. does not remove flex prefix)
	
	Params:
	matrix - SuperMatrix object
	filename - output filename
	out_encoding - output file encoding
	in_encoding - encoding of matrix features
	"""
	rows = matrix.NumberOfRows()
	rowfeatures = matrix.GetRowFeatures()
	file = codecs.open(filename, "w", out_encoding)

	file.write('word;tf;NZ;TF\n')
	for word in words:
		word_id = matrix.GetRowIndex(word)
		if word_id > rows: continue
		
		word_in_matrix = rowfeatures[word_id].GetLabel()
		word_in_matrix = unicode(word_in_matrix, matrix_encoding)
		
		nz = rowfeatures[word_id].GetNZ()
		TF= rowfeatures[word_id].GetGlobalTermFrequency()
		tf = rowfeatures[word_id].GetTermFrequency()
		
		line = "%s;%s;%s;%s\n" % (word_in_matrix,tf,nz,TF)
		file.write(line)
	file.close()

def extract_row_column_pairs(matrix, filename, output_encoding = 'cp1250' ,matrix_encoding = 'cp1250'):
	"""
	Extracts pair of (row,column) words from matrix for non zero cells.
	"""
	rows = matrix.NumberOfRows()
	rowfeatures = matrix.GetRowFeatures()
	columnfeatures = matrix.GetColumnFeatures()
	file = codecs.open(filename, "w", output_encoding)

	for row_no in xrange(rows):
		vals = matrix.GetRow(row_no).GetValues()
		for k in vals:
			word1 = rowfeatures[row_no].GetLabel()
			word2 = columnfeatures[k].GetLabel()
			if ':' in word1:
				word1 = word1[ word1.index(':')+1: ]
			if ':' in word2:
				word2 = word2[ word2.index(':')+1: ]
			if word1 == word2:
				continue
			line = "%s; %s ; %s\n" % (word1, word2, vals[k])
			file.write(line.decode(matrix_encoding))
	file.close()

def extract_single(bla, dirname, names):
	for name in names:
		filename = os.path.join(dirname, name)
		if os.path.isdir(filename):
			mdir,mname =  os.path.split(filename)
			if mname == 'polaczone': 
				continue
			matrix = pymatrices.LoadSuperMatrix(mdir, mname, False, True)
			extract_row_column_pairs(matrix, mname + '_pairs.txt', 'utf', 'utf')

def extract_dir(root):
	os.path.walk(root, extract_single, None)
	

if __name__ == '__main__':
	start_dir = '/mnt/data/workdir/macierze/supermatrix2/noun/relacje/'
	extract_dir(start_dir)
