#!/usr/bin/env python
# -*- coding: utf-8 -*-

#
# PyDotter
# ========
# PyDotter is a dot plot generator. It is designeed to plot multiple sequences against multiple
# sequences with the final graph having the originating sequences identified.

import sys, getopt
from Bio import SeqIO
from numpy import array, zeros, int32
import Image, ImageDraw, ImageChops, ImageOps
from time import time

# switch this to True to get debug messages
DEBUG=False

# default parameters
defimagesize=1024
defktup=12
defwindow=None				#autocalculate
defminmatch=12
defmismatch=0
defalpha=24

# version
versionstring="0.8.1"

# algorithm defines
LBDOT=0				#the default
ZANGYUANG=1

def usage():
	"""print the usage of the programme"""
	print "Usage:"
	print " %s [-h] -x xseq1 [-x xseq2 -x xseq3 ...] -y yseq1 [-y yseq2 -y yseq3 ...] -o outputfile [-s imagesize]"%sys.argv[0]
	print "-h\t--help\t\tthis helpful blurb"
	print "-x\t--xfile=\ta FASTA file for a sequence to be displayed on the x axis"
	print "-y\t--yfile=\ta FASTA file for a sequence to be displayed on the y axis"
	print "-C\t--conserved=\tspecify a fasta file to compare and mark conserved regions on the plot"
	print "-H\t--highlighe=\tspecify the colour and width of the conserved highlight markers. eg. red:2 or #f08033:3"
	print "-o\t--output=\tthe image file to write the output to."
	print "-s\t--size=\t\tthe output image size. Should be expressed as integer. Determines the longest side of the grid. real image will be slightly larger. [Default: %d]"%defimagesize
	print "-k\t--ktup=\t\tktuple size for tokenisation. [Default: %d]"%defktup
	print "-w\t--window=\twindow size for comparing mismatches. [Default: %s]"%((defwindow==None)and "autogenerated" or str(defwindow))
	print "-m\t--minmatch=\tminimum match length to be included. [Default: %d]"%defminmatch
	print "-d\t--mismatch=\tnumber of mismatches allowed per window of sequence. [Default: %d]"%defmismatch
	print "-S\t--save=\t\tsave the computed dotplot as the specified file"
	print "-L\t--load=\t\tload the precomputed dotplot from the specified file"
	print "-M\t--major=\toverride the automatic major tick seperation with this value"
	print "-T\t--minor=\toverride the automatic minor tick seperation with this value"
	print "-F\t--filter=\tfilter the dotplot to only include matches at least this long"
	print "-f\t--fine\tuse the slower finer algorithm to generate the dotplot. Works with lower ktuple and match sizes but takes significantly longer."
	print "-c\t--colour=\tspecify the colour to use for the sequence divisions. Specify as a word or a quoted hex colour string."
	print "-b\t--bound=\tspecify the colour to use for file bound division lines. Specify as a word or a quoted hex colour string."
	print "-a\t--alpha=\tspecify the value of alpha to use on drawing bounds. Value between 0 and 255. [Default: %d]"%defalpha
	print "-v\t--version\tdisplay version information"
	print
	print "Examples:"
	print
	print "calculate a dotplot for seq 1,2 and 3 vs seq 4 and 5. Output this to graph.png."
	print " %s -x seq1.fasta -x seq2.fasta -x seq3.fasta -y seq4.fasta -y seq5.fasta -o graph.png -s 1024"%sys.argv[0]
	print
	print "draw a dotplot and show conserved regions with two other sequences. save as conserved.png"
	print " %s -x seq1.fasta -y seq2.fasta -C seq3.fasta -C seq4.fasta -o conserved.png -s 1024"%sys.argv[0]
	print
	print "calculate a dotplot for a sequence versus itself. Save the outcome in a dotplot file for later rendering."
	print " %s -x sequence.fasta -y sequence.fasta -S sequence.dotplot"%sys.argv[0]
	print
	print "load a previously calculated dotplot and generate a large graph called image.jpg. Make the major ticks 10000 units apart."
	print " %s -L sequence.dotplot -s 4096 -o image.jpg -M 10000"%sys.argv[0]
	print
	print "load a previous dotplot and filter out all the matches less than 20 in length. Save the new dotplot."
	print " %s -L sequence.dotplot -F 20 -S sequence_majors.dotplot"%sys.argv[0]
	print
	print "load a previous dotplot and filter out all the matches less than 30 in length. Draw the resulting image as a graph. Use non standard colours. Save the output aswell"
	print " %s -L sequence.dotplot -F 30 -s 2048 -c red -b \"#008000\" -a 64 -o image.png -S sequence_over30.dotplot"%sys.argv[0]
	
def version():
	print "Freckle dot plotter version "+versionstring
	print "Copyright (C) 2007-2008 Center for Comparative Genomics, Murdoch University."
	print "Developed by Crispin Wellington <cwellington@ccg.murdoch.edu.au>"
	
def parseopts():
	"""parse the command line options and return them"""
	
	# default parameters
	xseq=[]
	yseq=[]
	conserved=[]
	outfile=None
	imagesize=defimagesize
	ktup=defktup
	window=defwindow
	minmatch=defminmatch
	mismatch=defmismatch
	savefile=None
	loadfile=None
	major=None
	minor=None
	filt=None
	seqbound=(0,0,255)
	filebound=(255,0,0)
	alpha=defalpha
	algo=LBDOT
	highlight=[(255,128,128),3]
	
	#our getopt definition strings
	shortopts="hx:y:o:s:k:w:m:d:S:L:M:T:F:vfc:b:a:C:H:"
	longopts=["help","xfile=","yfile=","output=","size=","ktup=","window=","minmatch=","mismatch=","save=","load=","major=","minor=","filter=","version","fine","colour=","bounds=","alpha=","conserved=","highlight="]
	
	if len(sys.argv[1:])==0:
		usage()
		sys.exit(0)
	
	try:
		opts,args=getopt.getopt(sys.argv[1:],shortopts,longopts)
	except getopt.GetoptError:
		usage()
		sys.exit(2)
		
	if len(args):
		print "ERROR: command line has extrenuous arguments"
		usage()
		sys.exit(1)
		
	for o,a in opts:
		if o in ("-h", "--help"):
			usage()
			sys.exit()
			
		elif o in ("-x","--xfile"):
			xseq.append(a)
			
		elif o in ("-y","--yfile"):
			yseq.append(a)

		elif o in ("-C","--conserved"):
			conserved.append(a)
			
		elif o in ("-H","--highlight"):
			highlight = [ parsecolour(a.split(':')[0]),int(a.split(':')[1])]
			
		elif o in ("-o","--output"):
			outfile=a
			
		elif o in ("-k","--ktup"):
			ktup=int(a)
			
		elif o in ("-w","--window"):
			window=int(a)
			
		elif o in ("-m","--minmatch"):
			minmatch=int(a)
			
		elif o in ("-d","--mismatch"):
			mismatch=int(a)
			
		elif o in ("-S","--save"):
			savefile=str(a)
			
		elif o in ("-L","--load"):
			loadfile=str(a)	
			
		elif o in ("-M","--major"):
			major=int(a)
			
		elif o in ("-T","--minor"):
			minor=int(a)
			
		elif o in ("-F","--filter"):
			filt=int(a)
		
		elif o in ("-f","--fine"):
			algo=ZANGYUANG
			
		elif o in ("-c","--colour"):
			seqbound=parsecolour(a)
		
		elif o in ("-b","--bounds"):
			filebound=parsecolour(a)
		
		elif o in ("-a","--alpha"):
			alpha=int(a)
			assert(alpha>=0 and alpha<=255)
			
		elif o in ("-v","--version"):
			version()
			sys.exit(0)
			
		elif o in ("-s","--size"):
			# try and parse the size string
			try:
				imagesize=int(a)
			except Exception, e:
				print "ERROR: cannot parse size string!",str(e)
				usage()
				sys.exit(3)
				
	#sanity check options
	if ktup>12:
		print "ERROR: maximum ktup size is 12"
		sys.exit(4)
	if ktup<4:
		print "ERROR: minimum ktup size is 12"
		sys.exit(4)
		
	if window != None and window<ktup:
		print "ERROR: window size must be at least ktup size"
		sys.exit(5)
	if minmatch<ktup:
		print "ERROR: minmatch size must be at least ktup size"
		sys.exit(6)
	if window != None and mismatch>=window:
		print "ERROR: mismatch size must be at less than window size, otherwise everything is a dot."
		sys.exit(7)
				
	return xseq, yseq, conserved, highlight, outfile, imagesize,ktup,window,minmatch,mismatch,savefile,loadfile,major,minor,filt,algo, seqbound, filebound, alpha
	
def parsecolour(colourstring):
	import string
	
	# first check english words	
	colourdb={
			'red':(255,0,0),
			'green':(0,255,0),
			'blue':(0,0,255),
			'yellow':(255,255,0),
			'cyan':(0,255,255),
			'magenta':(255,0,255)
		}
	if colourstring.lower() in colourdb.keys():
		return colourdb[colourstring.lower()]
	
	# next check hex hash style
	if colourstring[0]=='#' and len(colourstring)==7:
		assert( not False in [(ch>='0' and ch <= '9') or ( ch >='a' and ch <='f' ) for ch in colourstring[1:].lower()] )
		return tuple([string.atoi(colourstring[a:b],16) for a,b in [(1,3),(3,5),(5,7)]])
		
	raise Exception, "Unparsable colour string"
		

def main():
	xseqfiles,yseqfiles,conserved,highlight,outfile,imagesize,ktup,window,minmatch,mismatch,savefile,loadfile,major,minor,filt,algo, seqbound,filebound,alpha=parseopts()
	
	if DEBUG:
		print "xsequences:",xseqfiles
		print "ysequences:",yseqfiles
		print "conserved:",conserved
		print "highlight:",highlight
		print "outfile:",outfile
		print "imagesize",imagesize
		print "ktuple",ktup
		print "window",window
		print "minmatch",minmatch
		print "mismatch",mismatch
		print "loadfile",loadfile
		print "savefile",savefile
		print "major",major
		print "minor",minor
		print "filter",filt
	
	from DotPlot import DotPlot, LBDotPlot
	
	if algo==ZANGYUANG:
		plot=DotPlot(xseqfiles,yseqfiles,ktup, window, minmatch, mismatch)
	else:
		plot=LBDotPlot(xseqfiles,yseqfiles,ktup, window, minmatch, mismatch)
	
	if loadfile!=None:
		#load the dotstore from a previous run
		print "Loading dotplot",loadfile,"..."
		t=time()
		plot.Load(loadfile)
		print "done in",time()-t,"seconds"
	else:
		#calculate it
		print "Precalculating tables..."
		t=time()
		plot.CreateTables()
		print "done in",time()-t,"seconds"
	
		print "Calculating dotplot..."
		t=time()
		plot.CalculateDotStore()
		print "done in",time()-t,"seconds"

	if conserved:
		print "Calculating conserved regions in third sequence..."
		t=time()
		
		# we have some sequences to mark conserved regions from
		plot.CreateConservedStore()
		
		# TODO: handle more than one
		conservedlist=[]
		for num,cons in enumerate(conserved):
			print "processing conserved sequence",num+1,"..."
			conservedlist.append(plot.ProcessConservedRegions(cons))

		plot.CalculateConservedIntersection(conservedlist)
		print "done in",time()-t,"seconds"
		
		#sys.exit()
		#assert(dimension==0 or dimension==1)
		
		#start=self.ProcStart(start)
		#end=self.ProcEnd(dimension,end)
		
		## munge anything thats not 'ACGT' into a '.'
		#subseq=decodeseq(self.GetSubSequence(dimension,start,end))
		
		#table=(start,end,subseq,buildMappingTables(subseq, self.ktup))
		#self.tables[dimension][(start,end)]=table
		
		#return table
	
	#do we filter this?
	if filt!=None:
		print "Filtering dotplot..."
		t=time()
		plot.Filter(filt)
		print "done in",time()-t,"seconds"
	
	if savefile!=None:
		#save the dotstore
		print "Saving dotplot..."
		t=time()
		plot.Save(savefile)
		print "done in",time()-t,"seconds"
		
	
	xsize,ysize=plot.GetSequenceLength(0),plot.GetSequenceLength(1)
	
	# work out scale and thus final image width and height
	longest=(xsize>ysize) and xsize or ysize
	
	from math import ceil
	
	scale=ceil(float(longest)/float(imagesize))
	scale=float(longest)/float(imagesize)
	
	xoutput=float(xsize)/scale
	youtput=float(ysize)/scale
	
	if outfile!=None:
		plot.Interpolate()
		
		print "Indexing dotplot..."
		t=time()
		plot.IndexDotStores()
		print "done in",time()-t,"seconds"
		
		#print plot.dotstore
		#for key in plot.dotstore:
			#print key,":",plot.dotstore[key],"=>",len(plot.dotstore[key][0]),len(plot.dotstore[key][1])
		
		print "Calculating averaged grid..."
		t=time()
		plot.MakeAverageGrid(scale)
		print "done in",time()-t,"seconds"
		
		if conserved:
			print "Calculating conserved grid..."
			t=time()
			plot.MakeAverageGrid(scale, conserved=True)
			print "done in",time()-t,"seconds"
		
		print "Rendering image..."
		t=time()
		image=plot.MakeImage(major=major,minor=minor,seqbound=seqbound,filebound=filebound,alpha=alpha, invert=True)
		print "done in",time()-t,"seconds"
		
		if conserved:
			print "Rendering conserved image..."
			t=time()
			cimage=plot.MakeImage(major=major,minor=minor,seqbound=seqbound,filebound=filebound,alpha=alpha, conserved=True, invert=True)
			print "done in",time()-t,"seconds"
		
			COLOUR=highlight[0]
			WIDTH=highlight[1]
			
			#colourised = ImageOps.colorize(cimage, (255,255,255), tuple([255-COLOUR[i] for i in (0,1,2)]))
			colourised = ImageOps.colorize(cimage, (0,0,0), tuple([COLOUR[i] for i in (0,1,2)]))
			
			# make another image to paste this one onto
			newim = Image.new( "RGB", colourised.size, (0,0,0) )
			
			#grow_matrix=range(-WIDTH,WIDTH)
			
			draw = ImageDraw.Draw(newim)
			w,h = colourised.size
			for x in range(w):
				for y in range(h):
					pix = colourised.getpixel( (x,y) )
					if sum(pix):
						#[ newim.putpixel( (x+X,y+Y), tuple([255-a for a  in COLOUR]) ) for X in grow_matrix for Y in grow_matrix if X*Y<=WIDTH ]
						draw.ellipse( [(x-WIDTH,y-WIDTH),(x+WIDTH,y+WIDTH)], fill=tuple([255-a for a  in COLOUR]) )
			
			del draw
			
			colourised = newim
			
			image = image.convert("RGB")
			
			out = ImageChops.add( colourised, image )
			
			out=ImageChops.invert(out)
		
		print "Compiling image..."
		t=time()
		if not conserved:
			image = image.convert("RGB")
			out=ImageChops.invert(image)
			
		image2=plot.CompileImage(major=major,minor=minor,seqbound=seqbound,filebound=filebound,alpha=alpha, dotimage=out)
		print "done in",time()-t,"seconds"
		
		#cimage.
		
		image2.save(outfile)
		#colourised.save(outfile)
		
	sys.exit(0)


if __name__ == "__main__":
    main()
