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

# calculate Da and Rh from an ensemble of structures, given observed
# RDCs
#
usage = """

Calculate RDC alignment tensor using SVD given molecular structures

usage: calcTensor.py [options] <rdc.tbl> <pdb file1> [<pdb file2> ...]

 where rdc.tbl is one or more space- or colon- separated files containing
 XPLOR-SANI-style assignment statements, 
 and pdb file1, ... are one or more structure file which have previously
 been determined.

 options:
   -psf <psf filename>      - specify a psf filename. If omitted PSF info is
                              automatically generated from the pdb file.
   -showRDCs                - print out back-calculated RDC values
   -normType <normType>     - rdc normalization. One of NH, CH or none
                              [default value: NH]
   -useDistance             - include 1/rAB^3 distance dependence in rdc
                              calculation. By default this is only enabled
                              if a non-bonded rdc type is detected.
   -maxDa <value>           - set the maximum absolute value of Da
                              [default: 50].
   -fitSel <atom selection> - atom selection used to fit structures before
                              alignment tensor calculation.
   -rdcWeights <weights>    - colon- or space- separated list of weights used
                              to weight rdc experiments relative to each
                              other. [defaults to 1 for each].
   -crossValidate <files>   - colon- or space- separated list of files with
                              XPLOR SANI-style assignment statements of
                              measured RDC values which are not used for
                              fitting, but for which backcalculated RDCs
                              are printed with the calculated tensor. Use of
                              this term implies -showRDCs.
   -help-script             - print this message.
 
"""

(optList, args) = xplor.parseArguments(["psf:1",
                                        "useDistance:0",
                                        "maxDa:1",
                                        "showRDCs:0",
                                        "normType:1",
                                        "fitSel:1",
                                        "rdcWeights:1",
                                        "crossValidate:1", 
                                        "help-script:0",])


#
# selection used to fit structures before tensor calculation
#
rdcWeights=None
fitSel="name CA or name N or name C or name O"
printRDCs=False
normType='NH'
psf=None
useDistance=False
maxDa=50
crossTabs=[]   #cross-validation file names

for opt in optList:
    if opt[0]=="psf":
        psf=opt[1]
        pass
    if opt[0]=="useDistance":
        useDistance=True
        pass
    if opt[0]=="showRDCs":
        printRDCs=True
        pass
    if opt[0]=="maxDa":
        maxDa=float(opt[1])
        pass
    if opt[0]=="normType":
        normType=opt[1]
        pass
    if opt[0]=="fitSel":
        fitSel=opt[1]
        pass
    if opt[0]=='rdcWeights':
        if ':' in opt[1]:
            rdcWeights = opt[1].split(':')
        else:
            rdcWeights = opt[1].split()
            pass
        rdcWeights = map(lambda s: float(s), rdcWeights)
        pass
    if opt[0]=='crossValidate':
        printRDCs=True
        if ':' in opt[1]:
            crossTabs = opt[1].split(':')
        else:
            crossTabs = opt[1].split()
            pass
        pass
    if opt[0]=="help-script":
        print usage
        import sys
        sys.exit(0)
        pass
    pass

if len(args)<2:
    print usage
    sys.exit(1)
    pass

if ':' in args[0]:
    tables = args[0].split(':')
else:
    tables = args[0].split()
    pass
coordFiles = args[1:]

if not rdcWeights:
    rdcWeights=[1]*len(tables)
    pass

if len(rdcWeights)!=len(tables):
    raise Exception("there must be one weight per file")

import protocol

if not psf:
    from psfGen import pdbToPSF
    pdbToPSF(coordFiles[0])
else:
    protocol.initStruct(psf)
    pass

from varTensorTools import create_VarTensor, calcTensor
from rdcPotTools import create_RDCPot, scale_toNH, scale_toCH
from simulationTools import analyze

medium = create_VarTensor('medium')

if maxDa:
    medium.setDaMax(maxDa)
    pass


coordArray=[]
from atomSelAction import Fit
fitCoords=None
for file in coordFiles:
    protocol.initCoords(file)

    if not fitCoords:
        fitCoords = xplor.simulation.atomPosArr()
    else:
        if fitSel:
            AtomSel("all").apply(Fit(fitCoords,fitSel))
            pass
        pass

    coordArray.append( xplor.simulation.atomPosArr() )

    pass

from potList import PotList
rdcs=PotList()
cnt=0
for table in tables:
    rdcs.append( create_RDCPot(table,oTensor=medium,file=table) )
    if normType=='NH':
        scale_toNH( rdcs[cnt] )
    elif normType=='CH':
        scale_toCH( rdcs[cnt] )
    elif normType!='none':
        raise exception('normType must be one of: NH, CH or none')
    if useDistance:
        rdcs[cnt].setUseDistance(True)
        pass
    rdcs[cnt].setScale( rdcWeights[cnt] )
    rdcs[cnt].setShowAllRestraints( True )
    cnt += 1
    pass

crossTerms=PotList("cross")
cnt=0
for table in crossTabs:
    crossTerms.append( create_RDCPot("c_"+table,oTensor=medium,file=table) )
    if normType=='NH':
        scale_toNH( crossTerms[cnt] )
    elif normType=='CH':
        scale_toCH( crossTerms[cnt] )
    elif normType!='none':
        raise exception('normType must be one of: NH, CH or none')
    if useDistance:
        crossTerms[cnt].setUseDistance(True)
        pass
    crossTerms[cnt].setScale( rdcWeights[cnt] )
    crossTerms[cnt].setShowAllRestraints( True )
    cnt += 1
    pass


tensor = calcTensor(medium,coords=coordArray,expts=rdcs)


from vec3 import Vec3
#copy tensor atom positions
from atomSel import AtomSel
aniAtoms = AtomSel("resname ANI and resid %d" % medium.oAtom().residueNum())
aniPosAtoms = map(lambda a:(Vec3(a.pos()),a), aniAtoms)

#print analyze(medium)

print "Da: %7.2f    Rh: %7.3f" % (medium.Da() , medium.Rh())
print "Fit:"

for cnt in range(len(coordFiles)):
    xplor.simulation.setAtomPosArr(coordArray[cnt])
    #reset tensor atom positions
    for (pos,atom) in aniPosAtoms:
        atom.setPos(pos)
        pass
    print "  %-20s " % coordFiles[cnt],
    rdcs.calcEnergy()
    print "%20s: %7.3f" % (tables[0], rdcs[0].rms()),
    print
    pass

#print tensor
    
# uncomment the next line to get a printout of back-calculated rdcs
if printRDCs:
    print analyze(rdcs,crossTerms)
    pass