import tempfile,os,sys
from optparse import OptionParser
import shutil
from string import *
from os import listdir
from os.path import isfile, join
from plotting import *
import traceback
import subprocess

#Check if file exists and where it is. Return a full path to it
def check_file(filename):
	#See if it is in the current directory
	cwd = os.getcwd()
	try:
		with open(cwd+"/"+filename,'r'):
			return cwd+"/"+filename
	except IOError:
		#Not in the current directory, check if it is an absolute path then
		try:
			with open(filename,'r'):
				return filename
		except IOError:
			cleanup("File "+filename+" does not exist.")

#Create a temporary folder in the user directory - return it so that we can do whatever with it
def create_temp_folder():
	dirpath = tempfile.mkdtemp()
	return dirpath

#Remove all the temp directories that have been created in the course of execution
def cleanup(message):
	for folder in dumpster:
		if os.path.isdir(folder):
			print("Deleting",folder)
			shutil.rmtree(folder)
	print(message)	
	quit()
#Parse the options
usage = "USAGE: python PACEPipeline.py --abf ABFILE --abc ABCHAINS --agf ABFILE --agc ABCHAINS --o directory"
parser = OptionParser(usage=usage)

#Single file mode
parser.add_option("--light_seq",help="Antibody light chain sequence", dest="light_seq")
parser.add_option("--heavy_seq",help="Antibody heavy chain sequence", dest="heavy_seq")
parser.add_option("--ag_chain",help="Antigen chains to constrain, e.g. --c ABCD", dest="agc")
parser.add_option("--o",help="Output directory - where to write the results (full path)", dest="out_dir")
#user-defined epitope
parser.add_option("--epitope",help="User-defined epitope", dest="epitope")
#user-defined paratope
parser.add_option("--paratope",help="User-defined paratope", dest="paratope")
#Perform the modeling
parser.add_option("--modeling",help="If this is set, it will do modeling using a specific template", dest="modeling")

#List of temp folders that should be cleaned up after execution.
dumpster = []
local_path = os.path.dirname(os.path.realpath(__file__))

(options, args) = parser.parse_args()
output_dir = options.out_dir
if not os.path.isdir(options.out_dir):
	os.system("mkdir "+options.out_dir)

output_dir = options.out_dir

try:
	
	
	if (options.light_seq and options.heavy_seq and options.modeling and options.epitope and options.paratope and options.agc):
		
		#This will be modeled
		abfile = ""
		#This should have been submitted by the user
		agfile = output_dir+"/agstructure.pdb"
		
		#Firstly model the sequence.
		print("Modeling: ",options.modeling)
		print("Heavy - ",options.heavy_seq)
		print("Light - ",options.light_seq)	
		print("Epitope: ",options.paratope)
		print("Paratope: ",options.epitope)
		model_directory = output_dir+'/models'
		os.system('mkdir '+model_directory)
		
		s_file = open(output_dir+"/status.txt",'w')
		s_file.write("Modeling the antibody [1/4]")
		s_file.close()
		#On my server, I cannot do modeling.
		modeling_capable = True#You are model-capable in UCB
		if modeling_capable == True:
			if options.modeling=='CUSTOM':
			
				modeling_options = " --heavy "+options.heavy_seq+" --light "+options.light_seq+" --seqid 1.0 "
				for line in open(output_dir+'/userdefined_template.txt','r'):
					modeling_options+=" --template "+line.strip()+"_0 "
			
				print(modeling_options)
			
				#model_command = "python /pacedata/PACE_MAIN/ABDB_UI/ABDB/ModelAntibody/model_for_web.py "+modeling_options+" --outdir "+model_directory
				model_command = "python /home/icruser/public_html/csb/PACE/ABDB_UI/ABDB/ModelAntibody/model_for_web.py "+modeling_options+" --outdir "+model_directory
				os.system(model_command)
				print("Finished custom-template modeling")
				#Copy the model to the main execution directory
				model_files = [ f for f in listdir(model_directory) if isfile(join(model_directory,f)) ]
				print("Copying ",'cp -r '+model_directory+' '+output_dir)
				os.system('cp -r '+model_directory+' '+output_dir)
				for pot_model in model_files:
					if 'complete.pdb' in pot_model:
						os.system('cp '+join(model_directory,pot_model)+' '+output_dir+'/model.pdb')
						abfile = output_dir+'/model.pdb'
						break
				
			else:
			
				model_command = "modelantibody --seq "+options.heavy_seq+"/"+options.light_seq+" --outdir "+model_directory
				os.system(model_command)
				templates = [ f for f in listdir(model_directory) if not isfile(join(model_directory,f)) ]
				for template_folder in templates:
					if 'template00' in template_folder:
						model_files = [ f for f in listdir(join(model_directory,template_folder)) if isfile(join(model_directory,template_folder,f)) ]
						for pot_model in model_files:
							if 'complete.pdb' in pot_model:
								cp_cmd = 'cp '+join(model_directory,template_folder,pot_model)+' '+output_dir+'/model.pdb'
								os.system(cp_cmd)
								abfile = output_dir+'/model.pdb'
								print('Copying','cp -r '+model_directory+' '+output_dir)
								break
		
		#ZIP the results
		subprocess.call("zip -rq "+join(output_dir,'all_models.zip')+" models", cwd=output_dir, shell=True)
		abfile = output_dir+'/model.pdb'	
		#Input antibody file.
		abfile = check_file(abfile)
		print("Ab file submitted: ",abfile)
		#Input antigen file.
		agfile= check_file(agfile)
		print("Ag file submitted: ",agfile)
		
		s_file = open(output_dir+"/status.txt",'w')
		s_file.write("Dealing with the epitope and paratope [2/4]")
		s_file.close()
		os.mkdir(output_dir+"/output_ipatch")
		os.mkdir(output_dir+"/output_docking")
		#Run i-Patch - save results in its temp folder
		ipatch_temp = create_temp_folder()
		print("Ab i-Patch temp: ",ipatch_temp)	
		dumpster.append(ipatch_temp)
		ipatch_command = local_path+"/ABipatch/./ipatch_run_web.sh "+ipatch_temp+" "+local_path+"/ABipatch "+abfile+" "+agfile+" HL "+options.agc
		
		os.system(ipatch_command)
		
		os.system("cp -r "+ipatch_temp+"/* "+output_dir+"/output_ipatch")
		#Run EpiPredPipeline - save results in its temp folder
		docking_temp = create_temp_folder()
		print("Docking temp: ",docking_temp)	
		dumpster.append(docking_temp)
		userepi=""
		if (options.epitope=='CUSTOM'):
			options.user_epi = check_file(output_dir+'/userdefined_epitope.txt')
			userepi = "--user_epitope "+options.user_epi
		userpara = ""
		if options.paratope=='CUSTOM':
			options.user_epi = check_file(output_dir+'/userdefined_paratope.txt')
			userpara= '--user_paratope'
		docking_command = "python "+local_path+"/EpiPredPipeline/EpiPred.py --abf "+abfile+" --agf "+agfile+" --abc HL --agc "+options.agc+" --epitopes 3 --jobid "+docking_temp+" --para_ext 0 "+userepi+" "+userpara
		print(docking_command)
		s_file = open(output_dir+"/status.txt",'w')
		s_file.write("Docking Ab-Ag [3/4]")
		s_file.close()
		os.system(docking_command)
		os.system("cp -r "+docking_temp+"/* "+output_dir+"/output_docking")

		#Translate the docking results to chothia.
		trans_dir = output_dir+'/output_docking/decoys_translated'
		chothia_dir = output_dir+'/output_docking/'+'Chothia'
		os.mkdir(chothia_dir)
		onlyfiles = [ f for f in listdir(trans_dir) if isfile(join(trans_dir,f)) ]
		s_file = open(output_dir+"/status.txt",'w')
		s_file.write("Preparing the final output [4/4]")
		s_file.close()
		for decoy in onlyfiles:
			
			print("Chothia-translating ",decoy)
			tempfolder = create_temp_folder()
			command = "python "+local_path+"/EpiPredPipeline/Framer.py --f "+trans_dir+"/"+decoy+" --c HL --d chothia --o "+tempfolder
			os.system(command)
			os.system('cp '+tempfolder+'/chothia.pdb '+chothia_dir+'/'+decoy)
			os.system('rm -rf '+tempfolder)

		#Translate i-Patch results to Chothia
		tempfolder = create_temp_folder()
		command = "python "+local_path+"/EpiPredPipeline/Framer.py --f "+output_dir+"/output_ipatch/output_folder/solution.pdb --c HL --d chothia --o "+tempfolder
		print(command)
		os.system(command)
		os.system('cp '+tempfolder+'/chothia.pdb '+output_dir+'/output_ipatch/output_folder/chothia.pdb')
		os.system('rm -rf '+tempfolder)
		
		#Normalize results of docking and i-Patch
		print("java -jar "+local_path+"/DockContacts.jar "+output_dir+" 5.0")
		os.system("java -jar "+local_path+"/DockContacts.jar "+output_dir+" 5.0")
		print("java -jar "+local_path+"/Normalizer.jar "+output_dir+"/output_ipatch/output_folder/chothia.pdb "+output_dir+"/ABipatch_normalized.pdb")
		os.system("java -jar "+local_path+"/Normalizer.jar "+output_dir+"/output_ipatch/output_folder/chothia.pdb "+output_dir+"/ABipatch_normalized.pdb")
		translate_results(output_dir)#Called from plotter.py, writes out the CDR pngs
		#ZIP the results
		os.system('zip -jr '+join(output_dir,'all_decoys.zip')+' '+join(output_dir,'output_docking','top_decoys'))
	else:
		print("Something is missing from your input.")
		print(usage)
		
		
	s_file = open(output_dir+"/status.txt",'w')
	s_file.write("Done")
	s_file.close()
	cleanup("Program finished, see log for details")
except Exception as e:
	print("Failed, reason:",e)
	job_id = output_dir.split('/')
	print("================PROGRAM CRASHED=============================")
	print("Please contact sebastian.kelm@ucb.com citing this id: ",job_id[len(job_id)-2])
	s_file = open(output_dir+"/status.txt",'w')
	if e==None:#There is a bug -- it reaches the end ok but some invisible error pops up and the exceptino is none?
		s_file.write("Done")
	else:
		
		s_file.write("Crashed")
	s_file.close()
	cleanup("================PROGRAM CRASHED=============================")
	
