# -*- coding: utf-8 -*-
"""
Code needed for writing matrix created by tri2matrix using SuperMatrix format. 
For matrix description see documentation of write function. 

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

import os, codecs, sys

def write_matrix_content(lil, shape, filename):
	"""
	Writes lil-like matrix to supermatrix-crs format.
	A shape parameter is for convenience, as column based (supermatrx-ccs, i.e., matrix_ccs.mat)
	has row based description.
	"""
	with codecs.open(filename, 'w', 'utf8') as f:
		f.write('%d %d %d\n' % shape)
		rows = lil.shape[0]
		for row_no in xrange(rows):
			data = lil.data[row_no]
			ind = lil.rows[row_no]
			for p in zip(ind,data):
				f.write('%d %f ' % p )
			f.write('\n')

def write_features(features, filename):
	"""
	Writers features in supermatrix format.
	"""
	with codecs.open(filename, 'w', 'utf8') as f:
		f.write('label;nz;tf;TF;entropy\n')
		for fit in features:
			f.write('%s;%d;%d;%d;%f\n' % fit)

def write_labels(labels, filename):
	"""
	Writers labels of a matrix
	"""
	with codecs.open(filename, 'w', 'utf8') as f:
		for l in labels:
			f.write('%s\n' % l)

def write_info(matrix, filename):
	"""
	Writes info from a matrix. It uses two fields of a matrix object:
	- comment - a string description of matrix
	- tokens - number of tokens in a corpus used in building of the matrix
	"""
	with codecs.open(filename, 'w', 'utf8') as f:
		f.write('%s\n' % matrix.comment)
		f.write('%d' % matrix.tokens)
		
def write_udict(matrix, filename):
	"""
	Writes unigram dictionary to a file in legacy supermatrix format. I.e., it uses := as a separator.
	"""
	with codecs.open(filename, 'w', 'utf8') as f:
		for k in sorted(matrix.udict):
			f.write('%s := %d\n' % (k, matrix.udict[k]))
	
def write_boundary(matrix, filename):
	"""
	Writes boudaries of a matrix to a file.
	"""
	with codecs.open(filename, 'w', 'utf8') as f:
		for ind in matrix.boundaries:
			f.write('%s := %d\n' % (matrix.boundaries[ind], ind))

def write(matrix, dirname):
	"""
	Attempts to write matrix to given directory.

	Arguments:
	- matrix a lil_matrix like object with additional fields:
	  * row_features and column_features - a list of tuples (label,nz,tf,TF,entropy)
	  * comment - a string description of a matrix
	  * row_labels and column_labels - a lists of labels for rows and columns
	  * boundaries - a list of boundary indices for supermatrix
	  * name - a name for this matrix
	  * tokens - a number of tokens in the corpus used in creation of this matrix
	- dirname where a matrix will be written
	"""
	if not os.path.exists(dirname):
		os.makedirs(dirname)
	
	matrix_path = os.path.join(dirname, matrix.name)
	if os.path.exists(matrix_path):
		sys.stdout.write("Directory %s exists, will overwrite files\n" % matrix_path)
	else:
		os.makedirs(matrix_path)
	
	write_features(matrix.row_features, os.path.join(matrix_path, 'row_features.csv'))
	write_features(matrix.column_features, os.path.join(matrix_path, 'column_features.csv'))
	shape = (matrix.shape[0], matrix.shape[1], matrix.nnz)

	write_matrix_content(matrix, shape, os.path.join(matrix_path, 'matrix_crs.mat'))
	write_matrix_content(matrix.T, shape, os.path.join(matrix_path, 'matrix_ccs.mat'))

	write_info(matrix, os.path.join(matrix_path, 'info.txt'))
	write_udict(matrix, os.path.join(matrix_path, 'udict.txt'))

	write_labels(matrix.row_labels, os.path.join(matrix_path, matrix.name + '.rlabel'))
	write_labels(matrix.column_labels, os.path.join(matrix_path, matrix.name + '.clabel'))

	write_boundary(matrix, os.path.join(matrix_path, 'boundary.txt'))


