from vec3 import dot, cross, unitVec
from math import acos, pi
import trace

from dihedral import Dihedral

xplor.parseArguments() # check for typos on the command-line



def calcDihedral(a0,a1,a2,a3):
    " calc signed dihedral angle, in degrees"
    p0=a0.pos()
    p1=a1.pos()
    p2=a2.pos()
    p3=a3.pos()

    v1 = p0-p1
    v2 = p2-p1
    v3 = p1-p2
    v4 = p3-p2

    dp = dot(unitVec(cross(v1,v2)), unitVec(cross(v3,v4)))
    cp = cross(unitVec(cross(v1,v2)), unitVec(cross(v3,v4)))

    val = acos(dp)

    if (dot(cp,v2)<0) : val *= -1

    return val*180/pi




def getPhi(atomSel=0):
    """ determine all phi values in the given selection.
    Return as a dictionary whose keys are residue numbers.
    """
    if not atomSel: atomSel = "all"
    if type(atomSel) == type('string'):
        atomSel = AtomSel(atomSel)
        pass

    ret = {}
    for atomC in atomSel:
        if atomC.atomName() == 'C':
            resid = atomC.residueNum()+1
            try:
                atomN = AtomSel("resid %d and name N" % resid)[0]
                atomCA = AtomSel("resid %d and name CA" % resid)[0]
                atomCi = AtomSel("resid %d and name C" % resid)[0]
                
                ret[resid] = calcDihedral(atomC,atomN,atomCA,atomCi)
            except IndexError:
                pass
            pass
        pass
    return ret

def getPhi2(atomSel=0):
    """ determine all phi values in the given selection.
    Return as a dictionary whose keys are residue numbers.
    """
    if not atomSel: atomSel = "all"
    if type(atomSel) == type('string'):
        atomSel = AtomSel(atomSel)
        pass

    ret = {}
    for atomC in atomSel:
        if atomC.atomName() == 'C':
            resid = atomC.residueNum()+1
            try:
                phi = Dihedral(AtomSel("resid %d and name C" % (resid-1)),
                               AtomSel("resid %d and name N" % resid),
                               AtomSel("resid %d and name CA" % resid),
                               AtomSel("resid %d and name C" % resid)).value()
                
                ret[resid] = phi
            except IndexError:
                pass
            pass
        pass
    return ret

def getPsi(atomSel=0):
    """ determine all psi values in the given selection.
    Return as a dictionary whose keys are residue numbers.
    """
    if not atomSel: atomSel = "all"
    if type(atomSel) == type('string'):
        atomSel = AtomSel(atomSel)
        pass

    ret = {}
    for atomN in atomSel:
        if atomN.atomName() == 'N':
            resid = atomN.residueNum()
            try:
                atomCA = AtomSel("resid %d and name CA" % resid)[0]
                atomCi = AtomSel("resid %d and name C" % resid)[0]
                atomNp = AtomSel("resid %d and name N" % (resid+1))[0]
                
                ret[resid] = calcDihedral(atomN,atomCA,atomCi,atomNp)
            except IndexError:
                pass
            pass
        pass
    return ret

xplor.command("structure @gb3.psf end")
xplor.command("coor @gb3_refined_IV.pdb")

phiVals = getPhi("all")
phi2Vals = getPhi2("all")
psiVals = getPsi("all")
keys = phiVals.keys(); keys.sort()
for key in keys:
    try:
        phi = phiVals[key]
    except:
        phi=0
    try:
        psi = psiVals[key]
    except:
        phi=0
    try:
        phi2 = phi2Vals[key]
    except:
        psi2=0
    print "%3d %7.3f %7.3f " % (key,phi,psi)
    pass