#! /usr/bin/env python

description='''
DESCRIPTION	

	ABangle

	A tool to calculate and analyse the orientation between the VH and VL domains in antibodies.
	Command line interface.
	See README file for full description and examples.
'''

epilogue='''
AUTHOR
	
	2012
	James Dunbar - Oxford protein informatics group.
	james.dunbar@dtc.ox.ac.uk
	http://opig.stats.ox.ac.uk/
	http://www.stats.ox.ac.uk/~dunbar/
		
	Supervisors
	Prof Charlotte M Deane - Oxford protein informatics group.
	Dr Angelika Fuchs (Roche) and Dr Jiye Shi (UCB Celltech)
'''

import argparse, sys, os, time


if __name__=="__main__":
	parser = argparse.ArgumentParser(prog="ABangle", description=description, epilog=epilogue,formatter_class=argparse.RawDescriptionHelpFormatter)

	# For angle calculation for new structures from user
	parser.add_argument( '-i',type=str,nargs='+', default=[], help='File names of structures to domain orientations for. Can also be a ".dat" file containing a list of pdb files')
	parser.add_argument( '-store',choices=["y","n"],type=str,default=False,help='Store the data for user\'s structures?')
	parser.add_argument( '-online',action='store_true',default=False ,help='Use the ABnum public webserver')
	parser.add_argument( '-usernumbered',action='store_true',default=False ,help='Use the chothia numbering in the PDB file. Heavy and light chains must have H and L chain identifiers respectively')
	parser.add_argument( '-change_numbering_path',action='store_true',default=False ,help='Change the path to a local numbering program if available (e.g. abysis)')
	parser.add_argument( '-scfv',type=str,nargs='+', default=[], help='If the molecule is in single chain fv format, use this argument to tell the program which chain identifier it is (or they are)')

	parser.add_argument( '-change_data_path',action='store_true',default=False ,help='Change where to store user\'s data')
	parser.add_argument( '-clear',type=str,nargs='+',default=False ,help='Delete a structure from saved data. Use \'all\' to clear all')


	# For residue queries and finding structures
	parser.add_argument( '-r',type=str,nargs="+", default=[], help='Specify residue positions. e.g. L44 H22 L2.')   
	parser.add_argument( '-wa',type=str,nargs="+", default=[], help='With amino acid. eg. A S D')
	parser.add_argument( '-f',type=str,nargs="+", default=[], help='Factor query field. Choose from: HL, HC1, LC1, HC2, LC2, dc, Method, Res, Rfac, Bstate, AGtype, Hgroup, Lgroup, Ltype, Species, H3length.')
	parser.add_argument( '-wp',type=str,nargs="+", default=[], help='Factor with parameters. For class factors put the class parameter e.g. -f Species -wp Mouse. For numeric factor use e.g: -f HL -wp dx:X   for HL within dx of x. Or -f HL -wp gt:X for HL greater than X' )
	parser.add_argument( '-s',type=str,nargs='+', default=[], help='Request individual structures. These can either be the PDB code e.g. 12E8  or the code and specific chain identifiers e.g. 12E8_HL. If it has been pre-calculated it will be reported. Use "mine" to define the users saved structures or "all" to define all pre-calculated structures')
	parser.add_argument( '-target',type=str, default=False, help='Find Fvs which are structurally similar in terms of their angular measures to a target structure.')   

	# Search set selection
	parser.add_argument( '-useall', action='store_true', default=False, help="Use all the structures to search over. Not just a non-redundant set.")

	# Quiet
	parser.add_argument( '-q', action='store_true', help='Quiet. No headers are printed')

	# Show all the info
	parser.add_argument( '-showinfo', action='store_true', help='Show all the information for the structures in the set. Use -msa or -mr to see sequence')


	# For plotting measures
	parser.add_argument( '-plot',type=str,default=False, choices=["all","HL","HC1","HC2","LC1","LC2","dc"],help='Show a plot of the angle distributions of selected sets of structures against background distribution. Use \'all\' to plot all angles')
	parser.add_argument( '-lty',type=str,choices=['l','d'], default='', help='Show the sets using a density plot or vertical lines. Default makes a decision based on set size') 
	parser.add_argument( '-bg', type=str, default=False, help='File containing the data for backgound required for plots. Default is a non-redundant set.') 
	parser.add_argument( '-png',type=str,default=False, help='Save the angle plots as .png image file. No output to screen')

	# For showing structures in pymol
	parser.add_argument( '-pymol',action='store_true',help='Launch pymol and overlay selected structures. Structures are all aligned to the L chain core consensus strucuture.')
	parser.add_argument('-nalign', type=str, default='10', help='The maximum number of structures from each set to visualise in pymol. Use "all" to show all (Use with caution for large sets)')  
	parser.add_argument( '-useH',action='store_true',default=False,help='Align the structures to the VH consensus domain in pymol')
	parser.add_argument( '-showaxes',action='store_true',default=False,help='Show the axes in pymol')

	# For sequence or other information output
	parser.add_argument( '-msa',action='store_true',help='Output msa over Fv region')
	parser.add_argument( '-mr',type=str,nargs='+',default=False,help='Show msa of selected residues')
	parser.add_argument( '-seqid',action='store_true',help='Output the top 10 highest matched sequence identity pairs of Fvs from the first two sets defined (or within the first set if only one)')

	if len(sys.argv)==1:
	    parser.print_help() 
	    sys.exit(1)

	args=parser.parse_args()
	
	# Import the abangle modules
	from abangle import calculate, analyse 
	remove = []
	if args.pymol and not args.plot:
		args.plot = "all"

	try:
		#####################
		# System operations #
		#####################
		if args.change_data_path:
			calculate.ChangeUserPath(args)	
			reload( calculate )
			reload( analyse )

		if args.change_numbering_path or (args.i and not calculate.AnnotationProgPath and not args.online and not args.usernumbered):
			calculate.ChangeNumberingProgPath(args)	
			reload( calculate )
			reload( analyse )

		if args.clear:
			calculate.DeleteStore(args.clear)
			sys.exit()

		######################
		# Check dependencies #
		######################
	
		if args.plot:
			# Check that Rscript is available to show plotting. 
			if not analyse.which("Rscript"):
				sys.stderr.write( 'Rscript not found. Please download and install R and Rscript:\nhttp://www.r-project.org/\nPlotting unavailable, setting -plot option to False\n')
				args.plot=False
			
		if args.pymol:
			# Check that PyMOL is installed
			try:
				import pymol
			except ImportError:
				sys.stderr.write('Pymol not found. Please download and install from:\nhttp://pymol.org/\nVisualisation unavailable, setting -pymol option to False.\n')
				args.pymol = False	


		###############################
		# Calculating a new structure #
		###############################
		if args.i: 
			# Check TMalign is available
			if not analyse.which("TMalign"):
				raise Exception('Cannot find TMalign in path. Please install and ensure it is in your path.\nTMalign can be downloaded from:\n'\
				'http://zhanglab.ccmb.med.umich.edu/TM-align/\n'\
				'Reference: Y. Zhang and J. Skolnick, Nucl. Acids Res. 2005 33, 2302-9\n' ) 

			if len(args.i) == 1 and args.i[0].endswith(".dat"):
				# If a .dat file has been supplied expect a list of files - one per line
				args.i = map( str.strip, open( args.i[0] ).readlines() )

			# Locate Fvs and calculate the orientation measures for each.
			NewAngles = calculate.GetAngles(args)		

		else:
			NewAngles = []

		#######################################################################
		# Allow user to search ALL structures (not just in non-redundant set) #
		#######################################################################
		if args.useall:
			analyse.useall()

		##############################################################
		# Allow user to supply own background / subset of structures #
		##############################################################
		if args.bg:
			analyse.usebg( args )	

		##########################################################################
		# Find the amino acid combinations at specified set of residue positions # 	
		##########################################################################
		if analyse.CombinationSearch(args):
			sys.exit()

		#########################################################
		# Search for sets of structures with desired properties #
		#########################################################

		# Format from list to string
		args.wa = " ".join(args.wa)
		args.wp = " ".join(args.wp)

		# Find structurally similar Fvs to a target structure. Automates the angular parameter search
		if args.target:
			args = analyse.target(args)

		# If any query parameters have been defined then sets of structures that satisfy them are returned. Otherwise an empty list is returned.
		Sets=analyse.StructureSearch(args)		

		if args.s:
			Sets += analyse.UserSets( args.s )	

		#######################################
		# Display structures using R or pymol #
		#######################################

		if args.showinfo:
			analyse.OutputData(Sets, args)
		elif args.q:
			for s in Sets:
				for struc in s.structures:
					print struc
		else:
			for s in Sets:
				print s 
	
		if args.msa:
			args.mr = "all"
		if args.mr:
			analyse.ShowMsa( Sets, args )
		else:
			args.mr = "all"
		
		# Angle visualisation on plots 
		if args.plot:
			if not Sets:
				Sets += analyse.UserSets( ['all'] )
			remove=analyse.PlotAngles(Sets, args, NewAngles, show=True)
		elif args.pymol:
			# get the colours assigned to the sets
			remove=analyse.PlotAngles(Sets, args, NewAngles, show=False)

		# Structure visualisation using PyMOL
		if args.pymol and not args.i:
			analyse.ShowPymol(Sets, args)
		elif args.plot and not args.png:
			# Give display a chance to show the images before they are deleted.
			time.sleep(3)			
	
	except Exception, e:
		sys.stderr.write( "ABangle Error: %s\n"%str(e) )
	finally:
		# Delete temp files.
		for f in remove:
			try:
				os.remove( f )
			except:
				continue
		sys.exit()

