import sys
import argparse
import pylab
import numpy as np
import matplotlib.pyplot as plt
from heatmapcluster import heatmapcluster

def parseArgument():
	# Parse the input
	parser=argparse.ArgumentParser(description=\
			"Make a clustergram of a matrix")
	parser.add_argument("--matFileName", required=True,\
			help='File with matrix')
	parser.add_argument("--columnNamesFileName", required=True,\
			help='File with column names')
	parser.add_argument("--clusterFileName", required=True,\
			help='File where cluster information will be recorded')
	parser.add_argument("--metric", required=False, default='euclidean', \
			help='Metric for clustering')
	parser.add_argument("--contextIndex", required=False, type=int, \
			default=None, \
			help='Consider only a specific row')
	parser.add_argument("--logSignals", action='store_true', required=False,\
			help='log the signals')
	options = parser.parse_args()
	return options
	
def getStringList(stringFileName):
	# Get a list of strings from a file
	stringFile = open(stringFileName)
	stringList = [line.strip() for line in stringFile]
	stringFile.close()
	return stringList

def makeClustergram(options):
	# Make a clustergram of a matrix
	sys.setrecursionlimit(100000)
	mat = np.loadtxt(options.matFileName, dtype=np.float16)
	if options.logSignals:
		# log2 the signal values
		mat = np.log2(mat + 1)
	columnNames = getStringList(options.columnNamesFileName)
	if options.contextIndex is not None:
		# Considering the data in the context of a TF
		rowsToConsider = np.nonzero(mat[:,options.contextIndex])
		mat = mat[rowsToConsider,:].squeeze()
	rowNames = [str(i) for i in range(mat.shape[0])]
	pylab.figure()
	h = heatmapcluster(mat, rowNames, columnNames,
		num_row_clusters=None, num_col_clusters=None,
		label_fontsize=8,
		xlabel_rotation=-75,
		cmap=plt.cm.YlOrRd,
		show_colorbar=True,
		top_dendrogram=True,
		metric=options.metric)
	np.savetxt(options.clusterFileName, h.row_linkage, fmt='%.4f')
	pylab.show()

if __name__ == "__main__":
	options = parseArgument()
	makeClustergram(options)
