#!/usr/local/bin/python

# This is the GAG Binding Site Predictor that
# requires a PDB file as input where atomic
# charges have been placed in the B-factor column
# This can be obtained in many ways. One easy way
# to obtain such a file is to use psfgen from VMD
# to create a PSF and a PDB file and to replace the
# B-factor column in the PDB file by the charges in
# the PSF file.
# Created by Aurijit Sarkar, April 2014

import os, sys, math
from MolIO import *

unit_e = 1.6e-19

def GetDistance(iatom,atom):
	x1,y1,z1 = float(atom.x),float(atom.y),float(atom.z)
	x2,y2,z2 = float(iatom.x),float(iatom.y),float(iatom.z)
	dist = math.sqrt( (x1 - x2)*(x1 - x2) + (y1 - y2)*(y1 - y2) + (z1 - z2)*(z1 - z2) )
	return dist

def CalculateElectrostaticPotential(prot,atom):
	# atom must be a part of newmol. We will calculate the electrostatic potential due to all atoms of prot onto atom and return it
	sumPot = 0.0
	for iatom in prot.atoms:
		dist = GetDistance(iatom,atom)
		if dist > 0.0: # Do not attempt calculations on a zero distance. It causes a 'divide by zero' error. So, cannot calculate self potentials.
			q1 = float(iatom.bfactor)
			q2 = float(atom.bfactor)
			if q1 > 1 or q2 > 1:
				print 'too much charge. something is wrong',q1,q2
			delpot = ((unit_e*q1*unit_e*q2)/(dist*1e-10))
			sumPot = sumPot + delpot
	sumPot = (9e9*sumPot)
	#print sumPot
	kCalPerMol = sumPot*6.023e23*(1.0/4.184e3)
	return kCalPerMol

def IsNeutralDonor(atom):
	if atom.atomname.strip() not in ['N','ND2','SG','ND1','NE2','OG','OG1','OH']:
		return False
	else:
		if atom.resname == 'HSP':
			return False
		elif atom.atomname == 'NE2' and atom.resname == 'HSE':
			return False
		elif atom.atomname == 'ND1' and atom.resname == 'HSD':
			return False
		elif atom.resname in ['ARG','LYS']:
			return False
		else:
			return True
	return True

def CopyAtom(atom):
	newAtom = Atom()
	newAtom.set_attributes(atom.get_attributes())
	return newAtom

def IdentifyNeutralDonors(prot):
	newmol = Molecule()
	allNewAtoms = []
	for atom in prot.atoms:
		if IsNeutralDonor(atom):
			allNewAtoms.append(CopyAtom(atom))
	newmol.SetAtoms(allNewAtoms)
	return newmol

def CalculatePositiveChargeDensityAroundNeutralDonors(protf,op_prefix):
	protAtomDetails = ReadPDB(protf)
	allAtoms = GetAtomsFromAtomDetails(protAtomDetails)
	print 'total number of atoms found:',len(allAtoms)
	prot = GetProteinMolecule(allAtoms)
	
	newmol = IdentifyNeutralDonors(prot)
	print 'total number of neutral donors is:',len(newmol.atoms)
	
	for atom in newmol.atoms:
		potential = 0.0
		#print potential, atom.bfactor
		potential = CalculateElectrostaticPotential(prot,atom)
		#print potential, atom.bfactor
		atom.bfactor = str(round(potential,2))
	ofs = open(op_prefix+'_NeutralHBondDonor_GAGBindingPotentials.csv','w')
	
	ofs.write('AtomID,x,y,z,potential\n')
	for atom in newmol.atoms:
		name = '.'.join([atom.resname.strip(),atom.resnum.strip(),atom.atomname.strip()])
		line = ','.join([name,atom.x,atom.y,atom.z,atom.bfactor]) + '\n'
		ofs.write(line)
	ofs.close()

def main(argv):
	if len(argv) <> 3:
		print "Incorrect number of arguments for",argv[0]
		print "usage:",argv[0],"<Protein structure as PDB> <Output prefix>"
		sys.exit()
	
	prot_file = argv[1]
	op_prefix = argv[2]
	
	CalculatePositiveChargeDensityAroundNeutralDonors(prot_file,op_prefix)
	
	

if __name__ == "__main__":
	main(sys.argv)