#python
"""This script draws a set of overlayed histograms. It takes in a list of text files. Each text file contains data for one histogram and consists of N lines (vlaues separated by ' '). The first entry should be the labels. 

Input: 	text file listing files

Flags:	-h heading between  - DOESN'T wORK""
	-s save in the following folder
	-noline - no fit line
	-stack - produces stacked (and not overlayed histogram)
	- legend - add legend
	- nonorm - not normed (normed = 0)
	"""


import os 
import sys
import pylab
import numpy as np
import matplotlib.mlab as mlab

from statlib import stats as st


DATA = '/proj/aesop/data_lists'
PROC = '/proj/aesop/proc'

if 'bw' in sys.argv:
    COLOURS = ['grey', 'white']
else:
    COLOURS = ['blue', 'red', 'green', 'black', 'orange', 'magenta', 'yellow', 'purple', 'brown', 'beige', 'pink', 'grey', 'cyan']


def read_list(inf):
	fdic = {}
	names = []
	for line in (open(inf, 'r')):
		if line == '\n' or len(line) == 0:
			continue
		vals = line.split(' ')
		fdic[vals[0]] = read_file(open(vals[1].strip('\n'), 'r'))
		names.append(vals[0])
	return(fdic, names)

def read_file(file):
	data = {}
	names = []
	for line in file:
		if line == '\n' or len(line) == 0:
			continue
		vals = line.split(' ')
		data[vals[0]] = []
		names.append(vals[0])
		for v in vals[1:]:
			data[vals[0]].append(float(v))
	return (data, names)
	
	
def check_quotes(alist, string):
	for m in range(1, len(alist)):
		string = string+' '+alist[m]
		if '"' in alist[m]:
			return string

def read_arglist(arglist):
	save = 'no'
	h = ''
	fold = ''
	nl = 'yes'
	stack = 'no'
	leg = 'no'
	norm = 1
	inf = arglist[0]
	for n in range(1, len(arglist)):
		if arglist[n]  ==  '-h':
			h = check_quotes(arglist[n+1:], arglist[n+1])
		elif arglist[n]  ==  '-s':		
			save = 'yes'
			fold = arglist[n+1]
			if os.path.exists(fold) == False:
				os.mkdir(fold)
		elif arglist[n] == '-noline':
			nl = 'no'
		elif arglist[n] == '-stack':
			stack = 'yes'
		elif arglist[n] == '-legend':
			leg = 'yes'
		elif arglist[n] == '-nonorm':
			norm = 0	
	return (inf, h, save, fold, nl, stack, leg, norm)
		
def combine_arrays(data):
	xx = []
	handles = []
	labels = []
	i = 0
	a = 0
	for m in data.keys():
		if len(data[m]) == a:
			data[m].pop()
		x = np.array(data[m])
		xx.append(x)
		labels.append(m)
		handles.append(i)
		i+= 1
		a = len(data[m])
	return (xx, labels, handles)


def draw_hist(data, names, h, save, file, nl, stack, leg, norm):
	pylab.figure()
	pylab.subplots_adjust(hspace = 0.001)
	x_all = []
	for l in data.keys():
		print l, len(data[l][0])
		for m in data[l][0].keys():
			print m, len(data[l][0][m])
			x_all.extend(data[l][0][m])
	ax1 = pylab.subplot(len(data.keys()), 1, 1)
	n, bins, patches = pylab.hist(x_all, 50, normed = norm)
	axlim = ax1.get_xlim()
	bincenters  =  0.5*(bins[1:]+bins[:-1])
	y  =  mlab.normpdf( bincenters, np.mean(x_all), np.std(x_all))
	pylab.plot(bincenters, y, 'r--', linewidth = 1)
	pylab.ylabel('all')
	pylab.annotate('mean = %.2f' % (np.mean(x_all)), xy = (-50, 15), xycoords = 'axes points', fontsize = 8)
	pylab.annotate('std = %.2f' % (np.std(x_all)), xy = (-50, 5), xycoords = 'axes points', fontsize = 8)
	i = 1
	for l in names:
		handles = []
		labels = []
		print l, len(data[l][0])
		ax2 = pylab.subplot(len(data.keys()), 1, i, sharex = ax1)
		if stack == 'yes':
			(xx, labels, han) = combine_arrays(data[l][0])
			n, bins, patches = pylab.hist(xx, 60, normed = norm, histtype = 'barstacked', range = axlim)
			for ptch in patches:
				handles.append(ptch[0])
			if leg == 'yes':
				pylab.legend(handles, labels, ncol = 2)
			#print l, np.mean(xx)
		else:
			for m in data[l][0].keys():
				x = data[l][0][m]
				n, bins, patches = pylab.hist(x, 30, normed = norm, alpha = 0.7, histtype = 'bar')
				handles.append(patches[0])
				labels.append(m)
				if leg == 'yes':
					pylab.legend(handles, labels)
				bincenters  =  0.5*(bins[1:]+bins[:-1])
				y  =  mlab.normpdf( bincenters, np.mean(x), np.std(x))
				if nl == 'yes':
					pylab.plot(bincenters, y, '--', linewidth = 1)
		pylab.ylabel(l)
		pylab.yticks([], [])
		if not i == len(data.keys()):
			for tl in ax2.get_xticklabels():
				tl.set_visible(False)
		i += 1
	if save == 'yes':
		pylab.savefig(file)
	else:
		pylab.show()	
	
if __name__  ==  '__main__':
	arglist = sys.argv[1:]
	(inf, h, save, fold, nl, stack, leg, norm)  =  read_arglist(arglist)
	print 'Drawing histogram for files in %s, to be saved in %s' % (inf, fold)
	output = os.path.join(fold, '_'.join([os.path.basename(inf), 'histo'])+'.png')
	print output
	(data, names) = read_list(inf)
	draw_hist(data, names, h, save, output, nl, stack, leg, norm)	


