#!/usr/bin/env python
import argparse, sys, numpy, matplotlib, itertools, scipy
from matplotlib import pyplot
###########################################################
# Get command line arguments
def getParser():
	parser = argparse.ArgumentParser(description='Plot the values per genomic intervals.')
	parser.add_argument('-d',type=str,dest="dFiles",help="STR. Comma separated list of data files, all sorted by the same genomic intervales")
	parser.add_argument('-s',type=str,dest="species",help="STR. Species under study. It can be be mm10, hg38, or hg38_dm6")
	parser.add_argument('-c',type=str,dest="columns",help="STR. Comma separeted list of numbers indicating which columns contain the input in each data file. More than one column per data can be specified. For instance, '10|13,20|11' ")
	parser.add_argument('-o',type=str,dest="oFile",help="STR. Prefix of output png figures.")

	if len(sys.argv) == 1:
		print >> sys.stderr,parser.print_help()
		exit(0)
	return parser
#########################################################
def colormap(value,mu,L,U):
	if value>=mu:
		m=128.0/(U-mu)
		b=128.0*(U-2*mu)/(U-mu)
	else:
		m=128.0/(mu-L)
		b=128.0/(L-mu)*L
	y=int(round(m*value+b))
	return y
#########################################################
def main():
	args = getParser().parse_args()
	dataFiles = args.dFiles.strip().split(",")
	species = args.species
	prefix = args.oFile
	valCols = args.columns.strip().split(",")
	for i in range(len(valCols)):
		columns=[(int(col)-1) for col in valCols[i].split("|")]
		valCols[i]=columns
	##############
	# Chromomomal parameters
	chrom_sizes={}
	chrom_sizes["mm10"]={"chr1":195471971,"chr2":182113224,"chr3":160039680,"chr4":156508116,"chr5":151834684,"chr6":149736546,"chr7":145441459,"chr8":129401213,"chr9":124595110,"chr10":130694993,"chr11":122082543,"chr12":120129022,"chr13":120421639,"chr14":124902244,"chr15":104043685,"chr16":98207768,"chr17":94987271,"chr18":90702639,"chr19":61431566,"chrX":171031299,"chrY":91744698}
	chrom_sizes["hg38"]={"chr1":248956422,"chr2":242193529,"chr3":198295559,"chr4":190214555,"chr5":181538259,"chr6":170805979,"chr7":159345973,"chr8":145138636,"chr9":138394717,"chr10":133797422,"chr11":135086622,"chr12":133275309,"chr13":114364328,"chr14":107043718,"chr15":101991189,"chr16":90338345,"chr17":83257441,"chr18":80373285,"chr19":58617616,"chr20":64444167,"chr21":46709983,"chr22":50818468,"chrX":156040895,"chrY":57227415}
	chrom_sizes["hg38_dm6"]={"hg38_chr1":248956422,"hg38_chr2":242193529,"hg38_chr3":198295559,"hg38_chr4":190214555,"hg38_chr5":181538259,"hg38_chr6":170805979,"hg38_chr7":159345973,"hg38_chr8":145138636,"hg38_chr9":138394717,"hg38_chr10":133797422,"hg38_chr11":135086622,"hg38_chr12":133275309,"hg38_chr13":114364328,"hg38_chr14":107043718,"hg38_chr15":101991189,"hg38_chr16":90338345,"hg38_chr17":83257441,"hg38_chr18":80373285,"hg38_chr19":58617616,"hg38_chr20":64444167,"hg38_chr21":46709983,"hg38_chr22":50818468,"hg38_chrX":156040895,"hg38_chrY":57227415,"dm6_chr2R":25286936,"dm6_chr2L":23513712,"dm6_chr3R":32079331,"dm6_chr3L":28110227,"dm6_chr4":1348131,"dm6_chrX":23542271,"dm6_chrY":3667352}
	chrom_names={}
	chrom_names["mm10"]=numpy.array(["chr1","chr2","chr3","chr4","chr5","chr6","chr7","chr8","chr9","chr10","chr11","chr12","chr13","chr14","chr15","chr16","chr17","chr18","chr19","chrX","chrY"])
	chrom_names["hg38"]=numpy.array(["chr1","chr2","chr3","chr4","chr5","chr6","chr7","chr8","chr9","chr10","chr11","chr12","chr13","chr14","chr15","chr16","chr17","chr18","chr19","chr20","chr21","chr22","chrX","chrY"])
	chrom_names["hg38_dm6"]=numpy.array(["hg38_chr1","hg38_chr2","hg38_chr3","hg38_chr4","hg38_chr5","hg38_chr6","hg38_chr7","hg38_chr8","hg38_chr9","hg38_chr10","hg38_chr11","hg38_chr12","hg38_chr13","hg38_chr14","hg38_chr15","hg38_chr16","hg38_chr17","hg38_chr18","hg38_chr19","hg38_chr20","hg38_chr21","hg38_chr22","hg38_chrX","hg38_chrY","dm6_chr2R","dm6_chr2L","dm6_chr3R","dm6_chr3L","dm6_chr4","dm6_chrX","dm6_chrY"])
	
	chrom_names_short={}
	chrom_names_short["mm10"]=numpy.array(["1","2","3","4","5","6","7","8","9","10","11","12","13","14","15","16","17","18","19","X","Y"])
	chrom_names_short["hg38"]=numpy.array(["1","2","3","4","5","6","7","8","9","10","11","12","13","14","15","16","17","18","19","20","21","22","X","Y"])
	chrom_names_short["hg38_dm6"]=numpy.array(["1","2","3","4","5","6","7","8","9","10","11","12","13","14","15","16","17","18","19","20","21","22","X","Y","2R","2L","3R","3L","4","X","Y"])
	chrom_idxs={}
	chrom_idxs["hg38"]=numpy.array(range(24))
	chrom_idxs["hg38_dm6"]=numpy.array(range(31))
	chrom_idxs["mm10"]=numpy.array(range(21))
	chrom_pos={}
	chrom_pos["hg38_dm6"]=range(10,320,10)
	chrom_pos["hg38"]=range(10,250,10)
	chrom_pos["mm10"]=range(10,220,10)
	# Prepare figure
	#fig, ax = pyplot.subplots(1)
	cmap   = matplotlib.colors.LinearSegmentedColormap.from_list('my_colormap',['blue','white','red'],256)
	# Plot chromosome backgrounds
	for chrom_idx in chrom_idxs[species]:
		for i in range( len(dataFiles) ):
			widthx=10.0/len(dataFiles) - 1
			left=(chrom_pos[species][chrom_idx]-5) + (1 + widthx)*i
			#left=(chrom_pos[species][chrom_idx]-3)#+i*10+chrom_idx*10
			chrom_name=chrom_names[species][chrom_idx]
			#pyplot.broken_barh([(left,6)],(0,chrom_sizes[species][chrom_name]),facecolor="grey",alpha=0.3)
			pyplot.broken_barh([(left,widthx)],(0,chrom_sizes[species][chrom_name]),facecolor="grey",alpha=0.3)
	# Plot data
	for i,dataFile in enumerate(dataFiles):
		# Compute mu, L, and U
		values=numpy.array([])
		for n,line in enumerate(open(dataFile,"r")):
			if n==0: continue
			line=line.strip().split("\t")
			value=numpy.log10(sum([float(line[j]) for j in valCols[i] ])+1) # 16 normalized count
			values=numpy.append(values,value)
		mu=numpy.nanpercentile(values,50)
		L=numpy.nanpercentile(values,5)
		U=numpy.nanpercentile(values,95)
		# Fill chromosomes
		for n,line in enumerate(open(dataFile,"r")):
			if n==0: continue
			line=line.strip().split("\t")
			chrom_name=line[0]
			bottom=int(line[1])
			width=int(line[2])-bottom
			value=numpy.log10(sum([float(line[j]) for j in valCols[i] ])+1) # 16 normalized count

			chrom_idx=chrom_idxs[species][chrom_names[species]==chrom_name][0]

			widthx=10.0/len(dataFiles) - 1
			left=(chrom_pos[species][chrom_idx]-5) + (1 + widthx)*i
			#left=(chrom_pos[species][chrom_idx]-3)#+i*10+chrom_idx*10
			pyplot.broken_barh([(left,widthx)],(bottom,width),facecolor=cmap(colormap(value,mu,L,U)),alpha=1,edgecolor='none')
			#pyplot.broken_barh([(left,6)],(bottom,width),facecolor=cmap(colormap(value,mu,L,U)),alpha=1,edgecolor='none')
	# Set figure parameters
	pyplot.xlim(0,len(chrom_names_short[species])*10+7)
	pyplot.xticks(range(10,len(chrom_names_short[species])*10+10,10),chrom_names_short[species])
	pyplot.ylabel('Lengths [bases]')
	pyplot.xlabel('Chromosomes')
	pyplot.title(prefix)
	pyplot.savefig(prefix+".png")
########################################################
if __name__ == '__main__':
    main()
