import os,sys,random
from Bio.PDB import PDBParser
from Bio.PDB import PDBIO
import Bio.PDB
from optparse import OptionParser
from os import listdir
from os.path import isfile, join

def runPSA(code,the_rand):
	cwd = os.getcwd()
	dobavka = ""
	if (sys.platform=='darwin'):
		dobavka = "_mac"
	os.system(mypath+"/aux/psa"+dobavka+" -t "+mypath+"/temp"+the_rand+"/"+code+".pdb > "+mypath+"/temp"+the_rand+"/temp"+code+".txt")

#Single file mode - AB
def constrain(options,code,the_rand):
	#forward and backward mappings	
	mapping = dict()
	mapping["fwd"] = dict()
	mapping["bwd"] = dict()
	p = PDBParser()
	chains_to_choose = []
	if (code=='AB'):
		pdb_file = options.file_ab
		chains = options.chains_ab
	else:    
		pdb_file = options.file_ag
		chains = options.chains_ag

        structure = p.get_structure("input", pdb_file)
        
        for chain in chains:
                chains_to_choose.append(chain)
        i = 1
	for model in structure:
		for chain in model:
			if chain.id not in chains_to_choose:
				#What if there already is a chain A and we are not interested in it? 			
				if chain.id=='A':
					chain.id = '0'			
				continue
			else:
				for residue in chain:
					id = residue.id
					if id[0] == ' ':
						mapping['fwd'][i] = str(chain.id)+"_"+str(id[1])+"_"+str(id[2]).strip()
						mapping['bwd'][str(chain.id)+"_"+str(id[1])+"_"+str(id[2]).strip()] = i
						residue.id = (' ', i, ' ')
						i += 1
						
				chain.id = 'A'
	
        Select = Bio.PDB.Select
        class ConstrSelect(Select):
            def accept_chain(self, chain):
                #print dir(residue)
                
                if chain.id =='A':
                    return 1
                else:
                    return 0
	        
	w = PDBIO()
        w.set_structure(structure)
        w.save("TMP"+the_rand+".pdb",ConstrSelect())
        #Remove the HETATM and TER lines
	f_tmp = open("TMP"+the_rand+".pdb", 'r')
	f_out = open(mypath+"/temp"+the_rand+"/"+code+'.pdb', 'w')
	for line in f_tmp.readlines():
		if line[0:3]!="TER" and line[0:6]!="HETATM":
			f_out.write(line)
	f_tmp.close()
	f_out.close()	
	os.remove("TMP"+the_rand+".pdb")
	return mapping

#Returns true if the PDB has all the chains provided for the given pdb
def check_chains(options,code):
	#Single file mode - AB
	p = PDBParser()
	if (code=='AB'):
		pdb_file = options.file_ab
		chains = options.chains_ab
	else:    
		pdb_file = options.file_ag
		chains = options.chains_ag

        structure = p.get_structure("input", pdb_file)
        chains_to_choose = []
	for chain in chains:
                chains_to_choose.append(chain)
	
        chains_in_file = []
	for model in structure:
		for chain in model:
			chains_in_file.append(chain.id)
	for chain in chains_to_choose:
		if chain not in chains_in_file:
			return False
	return True	

#Single file mode - AB
def constrainBACK(pdb_file,pdb_file_out,the_map):
	#forward and backward mappings	
	f = open(pdb_file,'r')
	f_out = open(pdb_file_out,'w')
	output=""
	for line in f.readlines():
		line = line.strip()
		if "ATOM" in line:
			chain = line[21]
			res_num = int(line[23:28].strip())
			mapped = the_map[res_num]
			new_chain = mapped.split("_")[0]
			new_resnum = mapped.split("_")[1]
			while len(new_resnum)<3:
				new_resnum=" "+new_resnum
			if len(mapped.split("_")[2])==0:
				new_line = line[0:21]+new_chain+line[22:23]+new_resnum+line[26:len(line)]
			else:
				new_line = line[0:21]+new_chain+line[22:23]+new_resnum+mapped.split("_")[2]+line[27:len(line)]			
			output+=new_line+"\n"
	f_out.write(output)
	f_out.close()
	f.close()

def fileExists(filename):
	try:
		with open(filename): pass
	except IOError:
		return False
	return True


def crash(options):
	s_file = open(options.job_id+"/output_folder/status.txt",'w')
	s_file.write("Crashed")
	s_file.close()
	quit()


#Parse the options
usage = "USAGE: python RunABIpatch.py --abf ABFILE --abc ABCHAINS --agf AGFILE --agc AGCHAINS --jobid JOBID "
parser = OptionParser(usage=usage)

#Single file mode
parser.add_option("--abf",help="Antibody file location", dest="file_ab")
parser.add_option("--abc",help="Antibody chains to constrain, e.g. -c ABCD", dest="chains_ab")
parser.add_option("--agf",help="Antigen file location", dest="file_ag")
parser.add_option("--agc",help="Antigen chains to constrain, e.g. -c ABCD", dest="chains_ag")
parser.add_option("--paratope",help="Paratope file", dest="para_file")
parser.add_option("--jobid",help="Job id for this process", dest="job_id")

(options, args) = parser.parse_args()
cwd = os.getcwd()
mypath = os.path.dirname(os.path.realpath(__file__))
the_rand = str(int(random.random()*1000000000))

print("Everything appears to be in order<br>")

os.system("mkdir "+mypath+"/temp"+the_rand)
os.system("mkdir "+mypath+"/temp"+the_rand+"/out")
os.system("mkdir "+mypath+"/temp"+the_rand+"/temp_out")
os.system("mkdir "+mypath+"/output_folder")
os.system("chmod a+rwx -R "+mypath+"/temp*")
try:
	if (options.file_ab and options.file_ag and options.chains_ab and options.chains_ag and options.job_id):
		
		s_file = open(options.job_id+"/output_folder/status.txt",'w')
		s_file.write("Predicting")
		s_file.close()
		#Constrain the antibody file	
		ab_map = constrain(options,"AB",the_rand)
		#Constrain the antigen file
		ag_map = constrain(options,"AG",the_rand)
		#Find surface residues for antibody (AB) and antigen (AG) using PSA - doing it from java fails too often
		runPSA("AB",the_rand)
	        runPSA("AG",the_rand)
		#Run AB i-Patch
		os.system("java -jar ABiPatch.jar temp"+the_rand+"/AB.pdb temp"+the_rand+"/AG.pdb temp"+the_rand+"/tempAB.txt temp"+the_rand+"/tempAG.txt output_folder")        
		#Remove all the temp files
		os.system("rm -r "+mypath+"/temp*")
		os.system("rm mapping* const*")
		#Translate the single chain antibody returned by AB i-Patch to it's original coordinates and chains
			
		constrainBACK(mypath+"/output_folder/antibody.pdb",mypath+"/output_folder/solution.pdb",ab_map["fwd"])
		os.system("cp "+mypath+"/output_folder/solution.pdb "+mypath+"/output_folder/antibody.pdb")

		#Remove the single chain antibody file
		os.system("rm "+mypath+"/output_folder/antibody.pdb")
		s_file = open(options.job_id+"/output_folder/status.txt",'w')
		s_file.write("Done")
	else:
		print("Not enough input arguments supplied")
		print(usage)

except:#Catches all errors that might have arisen from executing this file
	s_file = open(options.job_id+"/output_folder/status.txt",'w')
	s_file.write("Crashed")
	s_file.close()


