#!/usr/bin/python
import sys
from math import sqrt, pi
from itertools import permutations
from math import acos,atan2

# see table 5.5 (p 132, v 4.6.3) in the gromacs manual
# for the different function types


#from
#http://stackoverflow.com/questions/1984799/cross-product-of-2-different-\
#vectors-in-python
def crossproduct(a, b):
	c = [a[1]*b[2] - a[2]*b[1],
		a[2]*b[0] - a[0]*b[2],
		a[0]*b[1] - a[1]*b[0]]
	return c
#end of code snippet

# mostly from 
# http://www.emoticode.net/python/calculate-angle-between-two-vectors.html
def dotproduct(a,b):
	return sum([a[i]*b[i] for i in range(len(a))])

def veclength(a):
	length=sqrt(a[0]**2+a[1]**2+a[2]**2)
	return length

def ange(a,b,la,lb,angle_unit):
	dp=dotproduct(a,b)
	costheta=dp/(la*lb)
	
	# to prevent math errors (imag)
	# the angles should be fine
	# may need to validate for dihedrals
	if costheta>1: 
		costheta=costheta-1
	elif costheta<-1:
		costheta=costheta+1

	angle=acos(costheta)
	if angle_unit=='deg':
		angle=angle*180/pi
	return angle
# end of code snippet

def diheder(a,b,c,angle_unit):
	dihedral=atan2(veclength(crossproduct(crossproduct(a,b),\
	crossproduct(b,c))), dotproduct(crossproduct(a,b),crossproduct(b,c)))
	if angle_unit=='deg':
		dihedral=dihedral*180/pi
	return dihedral

def readatoms(infile):
	positions=[]
	f=open(infile,'r')
	atomno=-2
	for line in f:
		atomno+=1
		if atomno >=1:
			position=filter(None,line.rstrip('\n').split(' '))
			if len(position)>3:
				positions+=[[position[0],int(atomno),\
				float(position[1]),float(position[2]),\
				float(position[3])]]
	return positions

def makebonds(positions,rcutoff,prevent_hhbonds):
	
	bonds=[]
	
	for firstatom in positions:
		for secondatom in positions:
			distance=round(sqrt((firstatom[2]-secondatom[2])**2\
+(firstatom[3]-secondatom[3])**2+(firstatom[4]-secondatom[4])**2)/10.0,5)
			xyz=[[firstatom[2],firstatom[3],firstatom[4]],\
[secondatom[2],secondatom[3],secondatom[4]]]
			
			if distance<=rcutoff and firstatom[1]!=secondatom[1]:
				if prevent_hhbonds and (firstatom[0][0:1]=='H' and\
				 secondatom[0][0:1]=='H'):
					 pass
				elif firstatom[1]<secondatom[1]:
					bonds+=[[firstatom[1],secondatom[1],\
distance,firstatom[0],secondatom[0],xyz[0],xyz[1]]]
				else:
					bonds+=[[secondatom[1],firstatom[1],\
distance,firstatom[0],secondatom[0],xyz[1],xyz[0]]]
	return bonds

def dedupe_bonds(bonds):
	
	newbonds=[]
	for olditem in bonds:
		dupe=False
		for newitem in newbonds:
			if newitem[0]==olditem[0] and newitem[1]==olditem[1]:
				dupe=True
				break;
		if dupe==False:
			newbonds+=[olditem]
	return(newbonds)

def genvec(a,b):
	vec=[b[0]-a[0],b[1]-a[1],b[2]-a[2]]
	return vec

def findangles(bonds,angle_unit):
	# for atoms 1,2,3 we can have the following situations
	# 1-2, 1-3
	# 1-2, 2-3
	# 1-3, 2-3
	# The indices are sorted so that the lower number is always first
	
	angles=[]
	for firstbond in bonds:
				
		for secondbond in bonds:
			if firstbond[0]==secondbond[0] and not (firstbond[1]==\
secondbond[1]): # 1-2, 1-3
				vec=[genvec(firstbond[6],firstbond[5])]
				vec+=[genvec(secondbond[6],secondbond[5])]
				angle=ange(vec[0],vec[1],firstbond[2]*10,\
secondbond[2]*10,angle_unit)
				angles+=[[firstbond[1],firstbond[0],\
secondbond[1],angle,firstbond[4],firstbond[3],secondbond[4],firstbond[6],\
firstbond[5],secondbond[6]]]
			
			if firstbond[0]==secondbond[1] and not (firstbond[1]==\
secondbond[1]): # 1-2, 3-1
				#this should never be relevant since we've sorted the atom numbers
				pass
			
			if firstbond[1]==secondbond[0] and not (firstbond[0]==\
secondbond[1]): # 1-2, 2-3
				vec=[genvec(firstbond[5],firstbond[6])]
				vec+=[genvec(secondbond[6],secondbond[5])]
				
#				print firstbond[3],firstbond[0],firstbond[4],firstbond[1],secondbond[4],secondbond[1]
				
				angle=ange(vec[0],vec[1],firstbond[2]*10,\
secondbond[2]*10,angle_unit)
				angles+=[[firstbond[0],firstbond[1],\
secondbond[1],angle,firstbond[3],firstbond[4],secondbond[4],firstbond[5],\
firstbond[6],secondbond[6]]]
			
			if firstbond[1]==secondbond[1] and not (firstbond[0]==\
secondbond[0]): # 1-3, 2-3
				vec=[genvec(firstbond[6],firstbond[5])]
				vec+=[genvec(secondbond[6],secondbond[5])]
				angle=ange(vec[0],vec[1],firstbond[2]*10,\
secondbond[2]*10,angle_unit)
				angles+=[[firstbond[0],firstbond[1],\
secondbond[0],angle,firstbond[3],firstbond[4],secondbond[3],firstbond[5],\
firstbond[6],secondbond[5]]]				
	return angles

def dedupe_angles(angles):
	dupe=False
	newangles=[]
	for item in angles:
		dupe=False
		for anotheritem in newangles:
			if item[0]==anotheritem[2] and (item[2]==anotheritem[0]\
 and item[1]==anotheritem[1]):
				dupe=True
				break
		if dupe==False:
			newangles+=[item]
	
	newerangles=[]
	dupe=False
	
	for item in newangles:
		dupe=False
		for anotheritem in newerangles:
			if item[2]==anotheritem[2] and (item[0]==anotheritem[1]\
 and item[1]==anotheritem[0]):
				dupe_item=anotheritem
				dupe=True
				break
		
		if dupe==False:
			newerangles+=[item]
		elif dupe==True:
			if dupe_item[3]>item[3]:
				pass;
			else:
				newerangles[len(newerangles)-1]=item

	newestangles=[]
	dupe=False
	
	for item in newerangles:
		dupe=False
		for anotheritem in newestangles:
			if (sorted(item[0:3]) == sorted (anotheritem[0:3])):
				dupe_item=anotheritem
				dupe=True
				break
		
		if dupe==False:
			newestangles+=[item]
		elif dupe==True:
			if dupe_item[3]>item[3]:
				pass;
			else:
				newestangles[len(newestangles)-1]=item
	return newestangles

def finddihedrals(angles,bonds,angle_unit):
	dihedrals=[]
	for item in angles:
		for anotheritem in bonds:
			if item[2]==anotheritem[0]:
				vec=[genvec(item[7],item[8])]
				vec+=[genvec(item[8],item[9])]
				vec+=[genvec(item[9],anotheritem[6])]
				dihedral=diheder(vec[0],vec[1],vec[2],\
angle_unit)
				dihedrals+=[ [ item[0],item[1],item[2],\
anotheritem[1], dihedral, item[4],item[5],item[6], anotheritem[4] ] ]

			if item[0]==anotheritem[0] and not item[1]==anotheritem[1]:
				vec=[genvec(anotheritem[6],item[7])]
				vec+=[genvec(item[7],item[8])]
				vec+=[genvec(item[8],item[9])]
				dihedral=diheder(vec[0],vec[1],vec[2],\
angle_unit)
				dihedrals+=[ [ anotheritem[1],item[0],item[1],\
item[2], dihedral, anotheritem[4],item[4],item[5], item[6] ] ]
	return dihedrals

def dedup_dihedrals(dihedrals):
	newdihedrals=[]
	
	for item in dihedrals:
		dupe=False
		for anotheritem in newdihedrals:
			rev=anotheritem[0:4]
			rev.reverse()
			if item[0:4] == rev:
				dupe=True
		if not dupe:
			newdihedrals+=[item]
	return newdihedrals

def print_bonds(bonds,func):
	print '[bonds]'
	
	for item in bonds:
            if 'H ' not in str(item[3])+'\t'+str(item[4]):
                    print str(item[3])+'\t'+str(item[4])+'\t'+\
                    str("%1.5f"% item[2])+'\t;'+str(item[0])+'\t'+str(item[1])
	return 0

def print_angles(angles,func):
	print '[ angles ]'
	for item in angles:
            if 'H' not in str(item[4])+'\t'+str(item[5])+'\t'+str(item[6]):
                    print str(item[4])+'\t'+str(item[5])+'\t'+str(item[6])+'\t'+\
        str("%3.3f" % item[3])+'\t;'\
		+str(item[0])+'\t'+str(item[1])+'\t'+str(item[2])
	return 0

def print_dihedrals(dihedrals,func):
	print "[ dihedrals ]"
	force='50000.00'
	mult='2'
	for item in dihedrals:
            if 'H' not in str(item[5])+'\t'+str(item[6])+'\t'+str(item[7])+'t'+str(item[8]):
                    print str(item[5])+'\t'+str(item[6])+'\t'+str(item[7])+'\t'+\
                    str(item[8])+'\t'+func+'\t'+str("%3.2f" % item[4])+'\t;'\
                    +str(item[0])+'\t'+str(item[1])+\
                    '\t'+str(item[2])+'\t'+str(item[3])
	return 0

if __name__ == "__main__":
	infile=sys.argv[1]
	rcutoff=float(sys.argv[2]) # in nm
	itemstoprint=int(sys.argv[3])
	
	angle_unit='deg' #{'rad'|'deg'}
	prevent_hhbonds=True # False is safer -- it prevents bonds between
	# atoms whose names start with H

	positions=readatoms(infile)
	
	if itemstoprint==1:
		bonds=makebonds(positions,rcutoff,prevent_hhbonds)
		bonds=dedupe_bonds(bonds)	
	
		print_bonds(bonds,'6')
	
	if itemstoprint==2:
		bonds=makebonds(positions,rcutoff,prevent_hhbonds)
		bonds=dedupe_bonds(bonds)	

		angles=findangles(bonds,angle_unit)
		angles=dedupe_angles(angles)
		print_angles(angles,'1')

	if itemstoprint==3:
		bonds=makebonds(positions,rcutoff,prevent_hhbonds)
		bonds=dedupe_bonds(bonds)	

		angles=findangles(bonds,angle_unit)
		angles=dedupe_angles(angles)

		dihedrals=finddihedrals(angles,bonds,angle_unit)
		dihedrals=dedup_dihedrals(dihedrals)
		print_dihedrals(dihedrals,'1')