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

import sys

(opts,files) = xplor.parseArguments(['rdc:1','rdcWeights:1',
                                     'csa:1',
                                     "maxDa:1",
                                     'weights:1',
                                     'svdTol:1',
                                     "help-script"])

usage="""
usage: %s [options] -rdc <rdc table> <structure files>

determine an ensemble of alignment tensors from an ensemble of structures,
such that the weighted sum if dipolar couplings corresponds with observed
values.

rdc table is a colon- or space- separated list of dipolar coupling assignment
tables relating to the same medium.

structure files is a list of pdbs to use in determining the tensors.

options:
  -weights <weights>    - a colon-separated list of weights to use on
                          structure ensemble members.
  -rdcWeights <weights> - a colon-separated list of weights used to
                          weight rdc experiments relative to each other.
  -csa <csa table>      - not yet supported.
  -maxDa <value>        - set the maximum absolute value of Da
                          [default: 50].
 -svdTol <value>       - discard singular values less than value times the
                          average of the singular values.
  -help-script          - print this message.

tensor details, predicted and observed dipolar couplings and rmsd values are
printed.
 
"""

ensembleSize=1
rdcTables=''
csaTables=''
weights=None
rdcWeights=None
maxDa=50
svdTol=0.01
for opt in opts:
    if opt[0]=='help-script':
        print usage
        import sys
        sys.exit(0)
        pass
    if opt[0]=="maxDa":
        maxDa=float(opt[1])
        pass
    if opt[0]=='rdc':
        if ':' in opt[1]:
            rdcTables = opt[1].split(':')
        else:
            rdcTables = opt[1].split()
            pass
        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]=='svdTol':
        svdTol=float(opt[1])
        pass
    if opt[0]=='csa':
        raise Exception("csas not yet supported")
    if opt[0]=='weights':
        if ':' in opt[1]:
            weights = opt[1].split(':')
        else:
            weigths = opt[1].split()
            pass
        weights = map(lambda s: float(s), weights)
        pass
    pass


pdbFile=files[0]

import psfGen
psfGen.pdbToPSF(pdbFile)

Ne=len(files)


if not weights:
    weights=[1]*Ne
    pass

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

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

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

rdcPairs=[]
for i in range(len(rdcTables)):
    rdcPairs.append((rdcTables[i],rdcWeights[i]))


from ensembleSimulation import EnsembleSimulation

esim = EnsembleSimulation("esim",Ne)

eIndex = esim.member().memberIndex()
esim.setWeights( weights )

print >> sys.stderr, 'reading ensemble of size %d with weights:' % Ne,
print >> sys.stderr, map(lambda i:esim.weight(i), range(esim.size()))

import protocol
protocol.initCoords(files[eIndex])

from rdcPotTools import create_RDCPot, scale_toNH

from varTensorTools import create_VarTensor, calcTensor_ensemble, calcTensor
oTensor = create_VarTensor("medium")

if maxDa:
    oTensor.setDaMax(maxDa)
    pass

from potList import PotList
rdcs=PotList()
for (rdcTable,rdcWeight) in rdcPairs:
    print >> sys.stderr, 'reading %s with weight %f' % (rdcTable, rdcWeight)
    rdc =  create_RDCPot(rdcTable,file=rdcTable,oTensor=oTensor) 
    rdc.setShowAllRestraints(1)
    scale_toNH( rdc )
    rdc.setScale( rdcWeight )
    rdc.setAveType("sum")

    rdcs.append(rdc)
    pass

tensors = calcTensor_ensemble(oTensor,
                              svdTolerance=svdTol)
#tensors = calcTensor(oTensor)

from simulationTools import analyze
print analyze(rdcs)
#print "total energy:", rdcs.calcEnergy().energy