#
#
#
 
"""get bond and angle from pdb files."""
__author__ = "Randall J. Radmer"
__version__ = "1.0"
  
 
import sys, os
import getopt
import fileinput
import math
#

degPerRad = 180/math.pi


def dVec(vA, vB):
    return [vA[0]-vB[0], vA[1]-vB[1], vA[2]-vB[2]] 

def mag(vA):
    return math.sqrt(vA[0]*vA[0]+vA[1]*vA[1]+vA[2]*vA[2])

def dot(vA, vB):
    return vA[0]*vB[0]+vA[1]*vB[1]+vA[2]*vB[2]

def dotNorm(vA, vB):
    return dot(vA, vB)/math.sqrt(dot(vA, vA)*dot(vB, vB))

def cross(vA, vB):
    return [vA[1]*vB[2]-vA[2]*vB[1],
            vA[2]*vB[0]-vA[0]*vB[2],
            vA[0]*vB[1]-vA[1]*vB[0]]

def crossNorm(vA, vB):
    vOut=cross(vA, vB)
    d = mag(vOut)
    vOut[0] /= d
    vOut[1] /= d
    vOut[2] /= d
    return vOut



def doIt(args_proper, atomName, delBond, delAngle, delDihedral):
    xyzOld=None
    xyzOldOld=None
    xyzOldOldOld=None
    atomOld=999999
    atomOldOld=999999
    atomOldOldOld=999999
    bondHist={}
    angleHist={}
    dihedralHist={}
    for line in fileinput.input(args_proper):
        if line.find("END")==0 or line.find("TER")==0:
            xyzOld=None
            xyzOldOld=None
            xyzOldOldOld=None
            atomOld=999999
            atomOldOld=999999
            atomOldOldOld=999999
            continue
        if line.find("ATOM")!=0 or len(line)<55:
            continue

        if atomName and atomName != line[11:16].strip():
            continue

        atom = int(line[6:11])
        xyz = (float(line[30:38]),
               float(line[38:46]),
               float(line[46:54]))

        if xyzOld:
            bondLength = mag(dVec(xyz, xyzOld))
            key = int(round(bondLength/delBond))
            if bondHist.has_key(key):
                bondHist[key]+=1
            else:
                bondHist[key]=1

            if xyzOldOld:
                cosA=dotNorm(dVec(xyzOldOld, xyzOld), dVec(xyz, xyzOld))
                cosA = max(cosA, -1)
                cosA = min(cosA,  1)
                angle = math.acos(cosA)*degPerRad
                key = int(round(angle/delAngle))
                if angleHist.has_key(key):
                    angleHist[key]+=1
                else:
                    angleHist[key]=1
                if xyzOldOldOld:
                    mVec=cross(dVec(xyzOldOldOld, xyzOldOld),
                               dVec(xyzOld,       xyzOldOld))
                    nVec=cross(dVec(xyzOld, xyzOldOld),
                               dVec(xyzOld, xyz))
                    cosA = dotNorm(mVec, nVec)
                    cosA = max(cosA, -1)
                    cosA = min(cosA,  1)
                    dihedral = math.acos(cosA)*degPerRad
                    if (dot(dVec(xyzOldOldOld, xyzOldOld), nVec)<0):
                        dihedral=-dihedral
                    key = int(round(dihedral/delDihedral))
                    if dihedralHist.has_key(key):
                        dihedralHist[key]+=1
                    else:
                        dihedralHist[key]=1
        
        if xyzOld:
            print "FOR %4d to %4d, dist = %4.2f" \
                  % (min(atomOldOldOld, atomOldOld, atomOld),
                     atom, bondLength),
            if xyzOldOld:
                print " angle = %5.1f" % angle,
                if xyzOldOldOld:
                    print " dihedral = %6.1f" % dihedral
                else:
                    print
            else:
                print

        atomOldOldOld=atomOldOld
        atomOldOld=atomOld
        atomOld=atom

        xyzOldOldOld=xyzOldOld
        xyzOldOld=xyzOld
        xyzOld=xyz

    keys=bondHist.keys()
    total=0
    for key in keys:
        total+=bondHist[key]
    for i in range(min(keys), max(keys), 1):
        if bondHist.has_key(i):
            value = bondHist[i]/float(total)
        else:
            value = 0.0
        print 'BOND       %6.3f %7.4f' % (i*delBond, value)

    keys=angleHist.keys()
    if len(keys)<1:
        print "No Angles"
        return True

    total=0
    for key in keys:
        total+=angleHist[key]
    for i in range(min(keys), max(keys), 1):
        if angleHist.has_key(i):
            value = angleHist[i]/float(total)
        else:
            value = 0.0
        print 'ANGLE    %6.1f   %8.5f' % (i*delAngle, value)

    keys=dihedralHist.keys()
    if len(keys)<1:
        print "No Dihedrals"
        return True

    total=0
    for key in keys:
        total+=dihedralHist[key]
    for i in range(min(keys), max(keys), 1):
        if dihedralHist.has_key(i):
            value = dihedralHist[i]/float(total)
        else:
            value = 0.0
        print 'DIHEDRAL %6.1f   %7.4f' % (i*delDihedral, value)

    return True

def parseCommandLine():
    opts, args_proper = getopt.getopt(sys.argv[1:], 'ha:B:A:D:')
    aValue = ''
    BValue = 0.01
    AValue = 1.0
    DValue = 10.0
    for option, parameter in opts:
        if option=='-h': usageError()
        if option=='-a': aValue = parameter
        if option=='-B': BValue = float(parameter)
        if option=='-A': AValue = float(parameter)
        if option=='-D': DValue = float(parameter)
    return (args_proper, aValue, BValue, AValue, DValue)

def main():
    args_proper, atomName, delBond, delAngle, delDihedral = parseCommandLine()

    doIt(args_proper, atomName, delBond, delAngle, delDihedral)

    return

def usageError():
    print 'usage: %s fileName1 [fileName2 [fileName3]]' \
         % os.path.basename(sys.argv[0])
    sys.exit(1)

if __name__=='__main__':
    main()