#!/usr/bin/env /opt/exp_soft/enmr/CIRMMP/xplor-nih/2.21/64/bin/pyXplor

# script to analyze and average protein torsion angles
#
# usage:
#   ./tang.py -psf=[psf file] [pdb files] >average.info
#
#

from sys import argv
from selectTools import minResid, maxResid
from simulationTools import StructureLoop
from dihedral import Dihedral
from pdbTool import PDBTool
import protocol

usage="""usage:
  torsionReport [-psf=psfFile] [pdb files] >average.info

 where [pdb files] is one or more structure files
 and information averaged over all the structures is send to stdout

 for each structure, a printout of all the torsion angles is saved to 
 a file named structN.angles

  """

(optList,files) = xplor.parseArguments(["psf:1", 
                                        "help-script:0",
                                        ])

psfRead=0
for opt in optList:
    if opt[0]=='psf':
        psf = opt[1]
        protocol.initStruct(psf)
        psfRead=1
        pass
    if opt[0]=="help-script":
        print usage
        import sys
        sys.exit(0)
        pass
    pass


#if not psfRead: protocol.initStruct("/home/schwitrs/clore/gb3/gb3.psf")

def getAngles(resType,resid,segid):
    segSel = 'and segid "%s"' % segid
    angles = [("phi",  ("name C  and resid %d %s" % (resid-1,segSel),
                        "name N  and resid %d %s" % (resid,segSel),
                        "name CA and resid %d %s" % (resid,segSel),
                        "name C  and resid %d %s" % (resid,segSel) )),
              ("psi",  ("name N  and resid %d %s" % (resid,segSel),
                        "name CA and resid %d %s" % (resid,segSel),
                        "name C  and resid %d %s" % (resid,segSel),
                        "name N  and resid %d %s" % (resid+1,segSel))),
              ("omega",("name CA and resid %d %s" % (resid,segSel),
                        "name C  and resid %d %s" % (resid,segSel),
                        "name N  and resid %d %s" % (resid+1,segSel),
                        "name CA and resid %d %s" % (resid+1,segSel)))]
    try:
        for (name,atoms) in scTorsionAtoms[resType]:
            angles.append((name,
                           ("name %s and resid %d %s" % (atoms[0],
                                                         resid,segSel),
                            "name %s and resid %d %s" % (atoms[1],
                                                         resid,segSel),
                            "name %s and resid %d %s" % (atoms[2],
                                                         resid,segSel),
                            "name %s and resid %d %s" % (atoms[3],
                                                             resid,segSel))))
            pass
        pass
    except KeyError:
        print "unsupported residue type:", resType
        pass
        
    return angles

scTorsionAtoms = {}
scTorsionAtoms['GLY'] = []
scTorsionAtoms['ALA'] = []
scTorsionAtoms['SER'] = [('CHI1',('N', 'CA', 'CB', 'OG')),
                         ('CHI2',('CA','CB', 'OG', 'HG'))]
scTorsionAtoms['THR'] = [('CHI1' , ('N' , 'CA', 'CB', 'OG1')),
                         ('CHI21', ('CA', 'CB', 'OG1', 'HG1'))]
scTorsionAtoms['LYS'] = [('CHI1', ('N',  'CA', 'CB', 'CG' )),
                         ('CHI2', ('CA', 'CB', 'CG', 'CD' )),
                         ('CHI3', ('CB', 'CG', 'CD', 'CE' )),
                         ('CHI4', ('CG', 'CD', 'CE', 'NZ' ))]
scTorsionAtoms['CYS'] =[('CHI1', ('N',  'CA', 'CB', 'SG')),
                        ('CHI2', ('CA', 'CB', 'SG', 'HG'))]
scTorsionAtoms['MET'] = [('CHI1', ('N',  'CA', 'CB', 'CG' )),
                         ('CHI2', ('CA', 'CB', 'CG', 'SD' )),
                         ('CHI3', ('CB', 'CG', 'SD', 'CE' ))]
scTorsionAtoms['VAL'] =[('CHI1',  ('N',  'CA', 'CB',  'CG1' ))]
scTorsionAtoms['ILE'] = [('CHI1',  ('N',  'CA',  'CB',  'CG1')),
                         ('CHI21', ('CA', 'CB',  'CG1', 'CD1'))]
scTorsionAtoms['LEU'] = [('CHI1',  ('N',  'CA',  'CB',  'CG'   )),
                         ('CHI2',  ('CA', 'CB',  'CG',  'CD1'  ))]
scTorsionAtoms['ASP'] = [('CHI1',  ('N',  'CA', 'CB',  'CG'  )),
                         ('CHI2',  ('CA', 'CB', 'CG',  'OD1' ))]
scTorsionAtoms['ASN'] = [('CHI1',  ('N',  'CA', 'CB',  'CG'  )),
                         ('CHI2',  ('CA', 'CB', 'CG',  'OD1' ))]
scTorsionAtoms['GLU'] = [('CHI1',  ('N',  'CA', 'CB',  'CG'  )),
                         ('CHI2',  ('CA', 'CB', 'CG',  'CD'  )),
                         ('CHI3',  ('CB', 'CG', 'CD',  'OE1' ))]
scTorsionAtoms['GLN'] = [('CHI1',  ('N',  'CA', 'CB',  'CG'  )),
                         ('CHI2',  ('CA', 'CB', 'CG',  'CD'  )),
                         ('CHI3',  ('CB', 'CG', 'CD',  'OE1' ))]
scTorsionAtoms['ARG'] = [('CHI1',  ('N',  'CA', 'CB',  'CG'  )),
                         ('CHI2',  ('CA', 'CB', 'CG',  'CD'  )),
                         ('CHI3',  ('CB', 'CG', 'CD',  'NE'  )),
                         ('CHI4',  ('CG', 'CD', 'NE',  'CZ'  ))]
scTorsionAtoms['PRO'] = []
scTorsionAtoms['HIS'] = [('CHI1',  ('N',  'CA', 'CB', 'CG'  )),
                         ('CHI2',  ('CA', 'CB', 'CG', 'ND1' ))]
scTorsionAtoms['PHE'] = [('CHI1',  ('N',  'CA', 'CB', 'CG' )),
                         ('CHI2',  ('CA', 'CB', 'CG', 'CD1'))]
scTorsionAtoms['TYR'] = [('CHI1',   ('N',   'CA', 'CB', 'CG' )),
                         ('CHI2',   ('CA',  'CB', 'CG', 'CD1')),
                         ('CHI6',   ('CE1', 'CZ', 'OH', 'HH' ))]
scTorsionAtoms['TRP'] = [('CHI1',  ('N',  'CA', 'CB', 'CG' )),
                         ('CHI2',  ('CA', 'CB', 'CG', 'CD1'))]

binData={}
#binData['CHI1'] = (('180',180), ('+60',60), ('-60',-60))
#binData['CHI2'] = (('180',180), ('+60',60), ('-60',-60))
#binData['CHI3'] = (('180',180), ('+60',60), ('-60',-60))
#binData['CHI4'] = (('180',180), ('+60',60), ('-60',-60))
#binData['CHI5'] = (('180',180), ('+60',60), ('-60',-60))

binData['PHE_CHI2'] = (('+90',90), ('-90',-90))
binData['TYR_CHI2'] = (('+90',90), ('-90',-90))
binData['TRP_CHI2'] = (('+90',90), ('-90',-90))
binData['HIS_CHI2'] = (('+90',90), ('-90',-90))
binData['ASP_CHI2'] = (('0',0), ('180',180))
binData['ASN_CHI2'] = (('0',0), ('180',180))
binData['GLU_CHI3'] = (('0',0), ('180',180))
binData['GLN_CHI3'] = (('0',0), ('180',180))

def getBranch(theta1,theta2):
    """return the branch number n such that n*2*PI+theta1-theta2 is
    minimal"""
    return round((theta2-theta1)/360)

def getBin(resType,taname,val):
    closest = ('single',0,val)
    bins=[]
    if taname.startswith('CHI'):
        bins = (('180',180), ('+60',60), ('-60',-60))
        pass
    try:
        bins = binData[resType+'_'+taname]
    except KeyError:
        pass
    
    if bins:
        closest = ('',361,361)
        for (name,mean) in bins:
            val += 360*getBranch(val,mean)
            dist=abs(val-mean)
            if dist<closest[1]:
                closest = (name,dist,val)
                pass
            #print 'branch',mean,val,getBranch(val,mean), closest
            pass
        pass
        
    return closest
        

    

angleBins={}
def binAngle(uid,resType,taName,val):
    global angleBins
    if not angleBins.has_key(uid):
        angleBins[uid] = {}
        pass
    (binName,binMean,branchVal) = getBin(resType,taName,val)
    while 1:
        try:
            tabin = angleBins[uid][taName]
        except KeyError:
            angleBins[uid][taName] = {}
            tabin = angleBins[uid][taName]
            pass
        try:
            tabin[binName].append(branchVal)
            break
        except KeyError:
            tabin[binName] = []
            pass
        pass
    return
            
        
from math import pi

def binAngles(filename):
    pdbFile = PDBTool(filename)
    pdbFile.read()
    afile = open(pdbFile.filename()+".angles",'w')

    #FIX: this assumes only a single segid
    for atom in AtomSel("tag"):
        resid = atom.residueNum()
        segid = atom.segmentName()
        resType = atom.residueName()
        uid = "%4s %4d %4s" % (segid, resid,resType)
        print >> afile, uid,
        for (aname,atomSel) in getAngles( resType , resid, segid):
            val=0
            try:
                val=Dihedral(*atomSel).value() * 180./pi
                print >> afile,  "%7.2f" % val,
                binAngle(uid,resType,aname,val)
            except IndexError:
                print >> afile, "[%5s]"%aname,
                pass
            pass
        print >> afile
        pass
    return

for filename in  files:
    binAngles(filename)
    pass


def average(l):
    ret=0.
    if not l: return ret
    
    for el in l:
        ret += el
        pass
    ret /= len(l)
    return ret

def deviation(l):
    from math import sqrt
    ret=0.
    if not l: return ret

    ave = average(l)

    for el in l:
        ret += (el-ave)**2
        pass
    ret /= len(l)
    return sqrt(ret)


print "residue        angle   bin    cnt    ave     dev"
for atom in AtomSel("tag"):
    resid = atom.residueNum()
    segid = atom.segmentName()
    resType = atom.residueName()
    uid = "%4s %4d %4s" % (segid, resid,resType)
    if not angleBins.has_key(uid):
        continue
    angles=angleBins[uid].keys()
    angles.sort()
    for angle in angles:
        bins = angleBins[uid][angle].keys()
        bins.sort()
        for bin in bins:
            print "%14s %5s %6s" % (uid, angle, bin),
            print "%4d" % len(angleBins[uid][angle][bin]),
            print "%7.2f" % average(angleBins[uid][angle][bin]),
            print "%7.2f" % deviation(angleBins[uid][angle][bin])
            pass
        pass
    pass