#!/usr/bin/python
#
#
'''
map value to beta field in corrosponding pdb

  - Usage: map2pdb <xyz-inputfile> <NBO charges> <optional NBO file>
  - input data format e.g.: C  1   -0.08588

 The optional NBO File can be used to calculate the difference 
 in charge. It will be mapped into the User field.
'''

import sys

__author__ = 'Christian Fufezan'
__url__ = 'http://www.fufezan.net/contact.php'
__version__ = 1.2

class ENTRY(object):
	global xyz
	def __init__(self,x,y,z,i,name):
		self.x = round(float(x),3)
		self.y = round(float(y),3)
		self.z = round(float(z),3)
		self.atype = str(name)+str(i)
		self.type = str(name)
		self.idx = int(i)
		self.beta = 7.777
		self.user = 7.777
	
	def addBeta(self,atomname,atomidx,charge):
		try:
			if self.idx == atomidx:
				self.beta = float(charge)
			else:
				print '<missmatch Beta> arry has',self.idx,' - expected',atomidx 
		except:
			print 'Woot'
		return
	
	def addUser(self,atomname,atomidx,charge):
		try:
			if self.idx == atomidx:
				self.user = 10*(self.beta - float(charge))
				#print self.idx,self.beta,charge,self.user
			else:
				print '<missmatch User> arry has',self.idx,' - expected',atomidx
		except:
			print 'Woot'
		return
			
	def iscomplete(self):
		if diffcrg != '':
			if self.beta == 7.777 or self.user == 7.777:
				return 0
			else:
				return 1
		else:
			if self.beta == 7.777:
				return 0
			else:
				return 1
			
def readin_xyz(file):
	coordinates=open(file).xreadlines()
	i = 0
	xyz = []
	for line in coordinates:
		if not line.startswith('#') and line.strip() != '':
			chopped = line.split()
			if len(chopped) == 4:
				tmp = ENTRY(chopped[1],chopped[2],chopped[3],i+1,chopped[0])
				i += 1
				xyz.append(tmp)
	return xyz

def bp():
	sys.exit(1)


if (__name__ == '__main__'):
	if (len(sys.argv) != 3 and len(sys.argv) != 4):
		print __doc__
		print __version__
		bp()
	else:
		inputxyz = str(sys.argv[1])
		inputcrg = str(sys.argv[2])
		if ( len(sys.argv) == 4):
			diffcrg = str(sys.argv[3])
		else:
			diffcrg = ''
		Output = {}
		xyz = []
		charges = []
		'''Reading xyz hem file ...'''
		xyz = readin_xyz(inputxyz)
		'''print 'Adding charges ... '''
		crg = open(inputcrg).xreadlines()
		for line in crg:
			if not line.startswith('#') and line.strip() != '':
				chopped = line.split()
				if len(chopped) == 3:
					xyz[int(chopped[1])-1].addBeta(chopped[0],int(chopped[1]),chopped[2])
		if diffcrg != '':
			diffcrg = open(str(sys.argv[3])).xreadlines()
			for line in diffcrg:
				if not line.startswith('#') and line.strip() != '':
					chopped = line.split()
					if len(chopped) == 3:
						xyz[int(chopped[1])-1].addUser(chopped[0],int(chopped[1]),chopped[2])
			
		print 'HEADER    Hem-Project CCNY 2007'
		print 'REMARK    generate by map2pdb. xyz file + charge '
		print 'REMARK    !! verify atom <-> charge correlation !!'
		
		for i in range(len(xyz)):
			if not xyz[i].iscomplete():
				print "<Error> Could not match charges for all atomes ..."
				bp()
			else:
				#print 'ATOM     60  CG1 ILE A   8       8.487  10.934  16.126  1.00 20.08           C'  
				print 'ATOM    %(j)3d %(atype)4s UKN A %(j)3d      %(x)6s  %(y)6s  %(z)6s %(user)5s %(crg)5s           %(e)s' % {
				'j'     : xyz[i].idx,
				'atype' : xyz[i].atype,
				'x'     : str('%(tmp).3f' % {'tmp' : float(xyz[i].x)}),
				'y'     : str('%(tmp).3f' % {'tmp' : float(xyz[i].y)}),
				'z'     : str('%(tmp).3f' % {'tmp' : float(xyz[i].z)}),
				'crg'   : str('%(tmp).2f' % {'tmp' : float(xyz[i].beta)}),
				'user'	: str('%(tmp).2f' % {'tmp' : float(xyz[i].user)}),
				'e'     : xyz[i].type
				}
		print 'END'
		bp()