"""high-level refinement tools
"""

from os import environ as env
from sys import modules
from potList import PotList
from pdbTool import PDBTool

#used when creating temporary files to avoid races in EnsembleSimulation
#calculations
fileSuffix=""

class StructureLoop:
    """ class which performs loop over structure calculations.
    Constructor: StructureLoop()
      arguments:
        numStructures       -    number of structures to calculate
        startStructure      -    id of the first structure to calculate
        structLoopAction    -    a user-defined function which takes
                                 one argument: an instance of this class.
        pdbTemplate         -    template string used to create a filename
                                 for the pdbFile method.

      if numStructures<0,   existing files matching pdbTemplate are processed,
      starting at startStructure, and stopping when a file does not exist.
      This mode of operation does not work with structure parallelism, but
      it does work with ensemble parallelism.

     There are additional arguments if you would like an average structure
     to be calculated.
       averageFilename      - name for output structure file. If not specified,
                              the average structure will not be calculated.
       averagePotList       - potential list to use for average structure
                              calculation.
       averageContext       - function to call to setup average minimization
                              run.  
       averageFitSel        - atom selection used to fit structures in
                              calculation of average structure [name CA].
       averageCompSel       - atom selection used for heavy atom rmsd
                              comparison metric [not hydro].
       averageTopFraction   - fraction of structures to use in calculation of
                              average structure. The structures are sorted
                              by energy, and the top fraction is retained.
                              [1 - all structures].
       s.averageRefineSteps - number of minimization steps to take during
                              regularization refinement of the average
                              structure [50]

    method: run():
       performs numStructures loop of action. In each pass, the coordinates
       of the current Simulation are reset to their initial values and the
       instance variable count is incremented. Also, the global random seed
       is incremented. If the current simulation is an EnsembleSimulation,
       the seed-setting takes this into account.

       
    """
    def __init__(s,numStructures=-1,startStructure=0,
                 structLoopAction=0,pdbTemplate="",
                 averageFilename="",
                 averagePotList=PotList(),
                 averageContext=lambda : 1,
                 averageFitSel="name CA",
                 averageCompSel="not hydro",
                 averageTopFraction=1,
                 averageRefineSteps=50     ):
        import sys
        from simulation import Simulation_currentSimulation
        s.numStructures=numStructures
        s.structLoopAction = structLoopAction
        s.pdbTemplate = pdbTemplate
        s.processID = int( env["XPLOR_PROCESS"] )
        proc = s.processID
        #if env.has_key("NUM_THREADS"):
        #    print "num_threads =", env["NUM_THREADS"]
        numProcs = int( env["XPLOR_NUM_PROCESSES"] )

        s.averageFilename  =averageFilename
        s.averagePotList   =averagePotList 
        s.averageContext   =averageContext 
        s.averageFitSel    =averageFitSel  
        s.averageCompSel   =averageCompSel 
        s.averageTopFraction=averageTopFraction
        s.averageRefineSteps=averageRefineSteps

        sim = Simulation_currentSimulation()
        s.esim = 0
        s.initCoords = sim.atomPosArr()
        s.skip=0
        s.count=0
        if ( sim.type() == "EnsembleSimulation"):
            from ensembleSimulation import EnsembleSimulation_currentSimulation
            s.esim = EnsembleSimulation_currentSimulation()
            s.skip = s.esim.member().memberIndex()
            global fileSuffix
            fileSuffix=str(s.skip)
            pass
        
        if numStructures>=0:
            # this logic used if number of jobs is greater than the desired 
            # number of structures
            numProcs = min(numProcs,numStructures)
            if proc >= numProcs:
                sys.exit()
                pass
            
            s.start = (proc     * numStructures) / numProcs + startStructure
            s.stop  = ((proc+1) * numStructures) / numProcs + startStructure
            
        else:
            # setup for processing existing files-
            if s.esim: s.skip=-2000
            s.start=startStructure
            s.stop=-1
            pass
        return

    def run(s):
        import simulationWorld
        from simulation import Simulation_currentSimulation
        from inspect import currentframe, getouterframes
        simWorld = simulationWorld.SimulationWorld_world()
        sim = Simulation_currentSimulation()
        initCoords = sim.atomPosArr()
        initSeed = simWorld.random.seed()
        
        s.count = s.start
        while 1 and s.structLoopAction:
            if s.count==s.stop: break
            if s.numStructures<0:
                #logic for processing existing files- and the number
                # of structures is not specified.
                #stop if the file doesn't exist.
                try:
                    import os
                    os.stat( s.makeFilename(s.pdbTemplate) )
                except:
                    break
                pass
            if simWorld.logLevel()!='none':
                print "StructureLoop: calculating structure %d" % s.count
            s.randomSeed = initSeed + s.count + s.skip*s.numStructures
            simWorld.setRandomSeed( s.randomSeed )
            sim.setAtomPosArr( initCoords )
            
            if type(s.structLoopAction) == type("string"):
                #structLoopInfo = s
                #exec( s.structLoopAction, vars( modules["__main__"] ), locals() )
                global_dict = getouterframes( currentframe() )[1][0].f_globals
                local_dict = getouterframes( currentframe() )[1][0].f_locals
                local_dict["structLoopInfo"] = s
                exec( s.structLoopAction, global_dict, local_dict )
                del local_dict["structLoopInfo"]
            else:
                s.structLoopAction(s)
                pass
            if s.esim: s.esim.barrier()
            s.count += 1
            pass

        if not s.averageFilename: return

        waitForProcesses()
    
        #FIX: should only be run by thread 0 of ensemble calc
        if s.processID==0:
            structs = []
            s.averageContext()
            for s.count in range(s.start,s.start+s.numStructures):
                #read files
                s.pdbFile().read()
                sim.sync()
                #calc energy
                try:
                    energy = s.averagePotList.calcEnergy().energy
                except AttributeError:
                    energy = 0
                    pass
                structs.append( (energy, s.pdbFile().filename(),
                                 sim.atomPosArr()) )
                pass
            #sort structures by energy
            structs.sort() #( lambda x,y: cmp(x[0],y[0]) )
            
            #take top goodFraction structures
            from math import ceil
            goodNum = int(ceil(len(structs) * s.averageTopFraction))
            structs = structs[0:goodNum]

            #  - perform average
            ( aveCoords,bFactors, remarks ) = \
              calcAverageStruct( map(lambda x: x[2], structs) ,
                                 fitSel=s.averageFitSel,
                                     compSel=s.averageCompSel)

            sim.setAtomPosArr( aveCoords )

            remarks += "Filenames:  (energy)\n"
            for (e,file,struct) in structs:
                remarks += "    %s    (%f)\n" % (file,e)
                pass
            remarks += '\n'

            #refine this structure using full potList
            minimizeRefine(s.averagePotList,
                           refineSteps=s.averageRefineSteps)
            
            from atomSel import AtomSel
            avePDB = PDBTool(s.averageFilename)
            for atom in AtomSel('known'):
                avePDB.setAux2(atom,bFactors[atom.index()])
            avePDB.addRemarks(remarks)
            pl = s.averagePotList
            avePDB.addRemarks(analyze(pl,
                                      outFilename=s.averageFilename+'.viols'))
            avePDB.write()
            pass
        pass
        waitForProcesses()
        return
    
    def pdbFile(s):
        """ return a PDBTool object whose filename is generated by
            makeFilename() from the pdbTemplate argument of the class
            constructor.
            """
        #a reference to the PDBTool objet is saved here to allow
        #access like info.pdbFile().filename()
        # if this is not done, the PDBFile object is reaped before the
        # filename() string is returned: garbage out
        s.pdbFileObj = PDBTool( s.filename() )
        return s.pdbFileObj
    def filename(s):
        """ return filename generated by makeFilename() from the pdbTemplate
        argument of the class constructor.
        """
        if not s.pdbTemplate:
            raise ("pdbFile: pdbTemplate must be specified in the\n\t" +
                   "StructureLoop constructor")
        return s.makeFilename(s.pdbTemplate)
    def makeFilename(s,template):
        """ create a filename given a template. In the template:
             the following substitutions are made:
                 SCRIPT    -> name of input script (without .py suffix)
                 STRUCTURE -> the structure count
                 MEMBER    -> the ensemble member index (0 is no ensembnle)
        """
        from simulation import Simulation_currentSimulation
        sim = Simulation_currentSimulation()
        memberIndex = 0
        if sim.type() == "EnsembleSimulation":
            from ensembleSimulation import EnsembleSimulation_currentSimulation
            esim=EnsembleSimulation_currentSimulation()
            memberIndex = esim.member().memberIndex()
            pass
        return genFilename(template,s.count,memberIndex)
    def analyze(s,potList,altPotList=PotList()):
        """print violation info to violations file
        and return summary information as a string
        """
        return analyze(potList,altPotList,
                       s.filename() + ".viols")
        return
    pass


def calcAverageStruct(structs,fitSel,compSel):
    """compute unregularize average structure given structs.
    The structures are first fit using fitSel, and analysis is performed
    using compSel.
    """

    from cdsVector import norm, sqrt
    from atomSelAction import Fit, RMSD
    from atomSel import AtomSel

    from simulation import Simulation_currentSimulation
    sim = Simulation_currentSimulation()
                                      
    fitTo = structs[0]
    sim.setAtomPosArr( fitTo )
    aveCoords = sim.atomPosArr() # must be separate copy
    var = norm(aveCoords)**2
    for struct in structs[1:]:
        sim.setAtomPosArr(struct)
        AtomSel("known").apply( Fit(fitTo,fitSel) )
        aveCoords += sim.atomPosArr()
        var += norm(sim.atomPosArr())**2
        pass
    
    aveCoords /= len(structs)
    var /= len(structs)
    bFactors = sqrt(var - norm(aveCoords)**2)

    fitTo=aveCoords
    aveRMSD=0.
    aveRMSDcomp=0.
    for struct in structs:
        sim.setAtomPosArr(struct)
        AtomSel("known").apply( Fit(fitTo,fitSel) )
                    
        comparer=RMSD(fitTo)
        AtomSel(fitSel).apply(comparer)
        aveRMSD += comparer.rmsd()
                    
        AtomSel(compSel).apply(comparer)
        aveRMSDcomp += comparer.rmsd()
        pass
    
    aveRMSD      /= len(structs)
    aveRMSDcomp /= len(structs)

    remarks  = "average structure over %d files\n" % len(structs)
    remarks += "fitted using atoms: %s\n" % fitSel
    remarks += "RMSD diff. for fitted atoms: %f\n" % aveRMSD
    remarks += "comparison atoms: %s\n" % compSel
    remarks += "RMSD diff. for comparison atoms: %f\n" % aveRMSDcomp
    remarks += "B array (last column) is rms diff from mean\n"
    return (aveCoords,bFactors, remarks)

def minimizeRefine(potList,
                   refineSteps=50,
                   xplorPots=['BOND','ANGL','IMPR'],
                   scaleMultiplier=0.001
                   ):
    """ simple refinement using gradient minimization in Cartesian coordinates.
    Some potential terms are held fixed during minimization, while others are
    scaled up from a smaller value progressively, during each round of
    minimization.

    refineSteps specifies how many rounds of minimization to perform.
    xplorPots are XPLOR terms which are to always be present, and for which
              the scale constant is held constant.
    scaleMultiplier specifies the initial value for the scale constant/
    """
    pots = flattenPotList(potList)

    #first get the tensor atoms in reasonable shape-
    # averaging will have scrambled them
    from varTensorTools import getVarTensors, calcTensor
    varTensors = getVarTensors(pots)

    for t in varTensors:
        calcTensor(t)
        pass


    # refine here
    #  remove bond, angle, impr, vdw terms- if they exist- use them as is.
    # if they don't exist, add them in with default scale values.
    #
    # for rest of terms, loop over them, with MultRamp running from .01 .. 1
    # times the nominal scale values.
    from potList import PotList
    minPots = PotList()
    hasReqdTerms = {}
    for p in xplorPots: hasReqdTerms[p] = 0
    rampedParams = []
    for pot in pots:
        reqdTerm=0
        for pType in xplorPots:
            if potType(pot) == pType:
                minPots.append(pot)
                hasReqdTerms[pType] = 1
                reqdTerm=1
                continue
            pass
        if reqdTerm: continue
        minPots.append( pot )
        rampedParams.append( MultRamp( scaleMultiplier*pot.scale(),
                                       pot.scale(),
                                       lambda val: pot.setScale(val) ))
        pass
    from xplorPot import XplorPot
    for pType in xplorPots:
        if not hasReqdTerms[pType]: minPots.append( XplorPot(ptype) )
        pass
    
    
    #Cartesian minimization
    from ivm import IVM
    minc = IVM()
    import protocol
    protocol.initMinimize(minc,
                          potList=minPots,
                          dEPred=10)
    minc.breakAllBondsIn("not resname ANI")
    minc.setBaseAtoms( "not resname ANI" )


    import varTensorTools
    for m in varTensors:
        m.setFreedom("varyDa, varyRh")       #allow all tensor parameters float
        varTensorTools.topologySetup(minc,m) #setup tensor topology
        pass

    if refineSteps>0:
        from simulationTools import AnnealIVM
        AnnealIVM(initTemp =0, finalTemp=0,
                  numSteps=refineSteps,
                  ivm=minc,
                  rampedParams = rampedParams).run()
        pass
                             
    # make sure final tensor is consistent with structure
    for t in varTensors:
        calcTensor(t)
        pass

    return

def potType(pot):
    """ return the potential's type.
    For most potential terms, the type is given by the potName() accessor,
    but there are exceptions:
      for <m xplorPot>XplorPot, the returned type is given by instanceName()
      for <m avePot>AvePot, potType(subPot()) is returned.
    """
    if pot.potName()=='XplorPot': return pot.instanceName()
    if pot.potName()=='AvePot':   return potType(pot.subPot())
    return pot.potName()
    


def genFilename(template,
                structure=0,
                member=0):
        """ create a filename given a template. In the template:
             the following substitutions are made:
                 SCRIPT    -> name of input script (without .py suffix)
                 STRUCTURE -> the structure argument
                 MEMBER    -> the member argument
        """
        import re

        from sys import argv

        try:
            scriptName =  re.sub(r'\.[^.]*$','',argv[0])
        except IndexError:
            scriptName = "stdin"
            pass
        if scriptName.endswith(__name__) or scriptName=='':
            scriptName = "stdin"
            pass
        filename = template.replace("SCRIPT","%s" %scriptName)

        filename = filename.replace("STRUCTURE","%d" %structure)
        filename = filename.replace("MEMBER","%d" % member)
        return filename
    

def waitForProcesses():
    """wait for all processing calculating structures in parallel to complete
    their current computation.
    This routine acts like a barrier, suspending computation until all
    processes have reached this point.

    this routine assumes write access to the local directory, and that this
    directory is shared by all processes.
    """
    pollInterval = 5 #secs to wait between polling for process completion
    numProcs = int( env["XPLOR_NUM_PROCESSES"] )
    if numProcs==1:
        return

    processID = int( env["XPLOR_PROCESS"] )
    ppid = int( env["XPLOR_PPID"] )

    import re
    from sys import argv
    try:
        scriptName =  re.sub(r'\.[^.]*$','',argv[0])
    except IndexError:
        scriptName = "stdin"
        pass
    if scriptName.endswith(__name__):
        scriptName = "stdin"
        pass

    barrierFilePref = ".%s-barrier-%d-" % (scriptName,ppid)
    barrierFile = barrierFilePref+"%s" % processID

    #remove file if it already exists
    import os, time
    try:
#        os.stat(barrierFile)
        while 1:
            if open(barrierFile).read() == "run.\n": break
            time.sleep(pollInterval)
            pass
#        print "waitForProcesses: WARNING: barrier file already exists!"
#        os.unlink(barrierFile)
    except IOError:
        pass

    file = open(barrierFile,"w"); file.write("waiting.\n"); file.close()

    if processID==0:
        #wait until everyone is done.
        for cnt in range(1,numProcs):
            while 1:
                file = barrierFilePref+"%s" % cnt
                try:
                    print "trying ", file
                    if open(file).read()=="waiting.\n": break
                    pass
                except IOError:
                    pass
                print "waiting for process %d" % cnt
                time.sleep(pollInterval)
                pass
            pass
        print "waking all processes"
        #wake all processes
        for cnt in range(1,numProcs):
            file = barrierFilePref+"%s" % cnt
            open(file,"w").write("wake.\n")
            pass
        #wait til process is running again, then delete its barrier file
        for cnt in range(1,numProcs):
            file = barrierFilePref+"%s" % cnt
            while open(file).read() != "waking.\n":
                time.sleep(1)
                pass
            open(file,'w').write('run.\n')
            os.unlink(file)
            pass
        open(barrierFile,'w').write('run.\n')
        os.unlink(barrierFile)
    else: # processID!=0
        #wait until awakened by process 0
        while 1:
            if open(barrierFile).read() == "wake.\n":
                break
            time.sleep(pollInterval)
            pass
        #tell process 0 that we're awake.
        print "process awoken"
        open(barrierFile,"w").write("waking.\n")
        pass
    
        
    return

class RampedParameter:
    """Base class for ramped parameters.
      update() - increments value. It will not change the value beyond
                 that specified by stopValue
      value()  - return the current value
      init(ns)        - set number of steps and initialize val to startValue
      finalize()      - set value to the final value.
    """
    def __init__(s,action):
        from inspect import currentframe, getouterframes
        s.action = action
        s.callingFrame = getouterframes( currentframe() )[2][0]
        s.val = 0
        s.startValue = 0
        s.stopValue  = 0
        return
    def init(s,numSteps,caller=0):
        s.setNumSteps(numSteps)
        s.val = s.startValue
        s.runAction(caller)
        return
    def finalize(s,caller=0):
        s.val = s.stopValue
        s.runAction(caller)
        return
    def runAction(s,caller):
        if not s.action: return
        if type(s.action) == type("string"):
            global_dict = s.callingFrame.f_globals
            local_dict = s.callingFrame.f_locals
            local_dict["ParameterRampInfo"] = s
            local_dict["caller"] = caller
            exec( s.action.replace("VALUE","%f" %s.val),
                  global_dict, local_dict )
            del local_dict["ParameterRampInfo"]
        else:
            s.action(s.val)
            pass
        return
    def setNumSteps(s,numSteps):
        return
    def value(s):
        return s.val
    def update(s,caller=0):
        s.runAction(caller)
        return 0.
    pass


class MultRamp(RampedParameter):
    """convenience class for multiplicatively (geometrically)
    ramping a value from startValue to stopValue over numberSteps
    constructor: MultRamp(startValue, stopValue, action)
    methods:
      update() - increments value. It will not change the value beyond
                 that specified by stopValue
      value()  - return the current value
      setNumSteps(ns) - set number of steps
      init(ns)        - set number of steps and initialize val to startValue
      finalize()      - set value to the final value.
    """
    def __init__(s,startValue,stopValue,action=None):
        RampedParameter.__init__(s,action)

        s.val = startValue
        s.startValue = startValue
        s.stopValue  = stopValue
        s.dirIncreasing=0
        if stopValue>startValue: s.dirIncreasing=1
        return
    def update(s,caller=0):
        s.val *= s.factor
        if (s.dirIncreasing and s.val>s.stopValue) or \
           (not s.dirIncreasing and s.val<s.stopValue):
            s.val = s.stopValue
            pass
        s.runAction(caller)
            
        return s.val
    def setNumSteps(s,numSteps):
        s.factor=1
        if numSteps: s.factor = (s.stopValue/s.startValue)**(1./numSteps)
        return
    pass

class LinRamp(RampedParameter):
    """convenience class for linearly
    ramping a value from startValue to stopValue over numberSteps
    constructor: MultRamp(startValue, stopValue, action)
    methods:
      update() - increments value. It will not change the value beyond
                 that specified by stopValue
      value()  - return the current value
      setNumSteps(ns) - set number of steps
      init(ns)        - set number of steps and initialize val to startValue
    """
    def __init__(s,startValue,stopValue,action=None):
        RampedParameter.__init__(s,action)

        s.val = startValue
        s.startValue = startValue
        s.stopValue  = stopValue
        s.dirIncreasing=0
        if stopValue>startValue: s.dirIncreasing=1
        return
    def update(s,caller=0):
        s.val += s.factor
        if (s.dirIncreasing and s.val>s.stopValue) or \
           (not s.dirIncreasing and s.val<s.stopValue):
            s.val = s.stopValue
            pass
        s.runAction(caller)
            
        return s.val
    def setNumSteps(s,numSteps):
        s.factor=0
        if numSteps>0: s.factor = (s.stopValue-s.startValue)/numSteps
        return
    pass

class StaticRamp(RampedParameter):
    """convenience class for static parameter setup.
      update() - increments value. It will not change the value beyond
                 that specified by stopValue
      value()  - return the current value
      setNumSteps(ns) - set number of steps
      init(ns)        - set number of steps and initialize val to startValue
    """
    def __init__(s,action):
        RampedParameter.__init__(s,action)
        return
    def update(s,caller=0):
        s.runAction(caller)
            
        return s.val
    pass

class InitialParams:
    """constructor takes a list of ramped parameters. The constructor invokes
    each parameter such that it set to its initial value.
    Also, this object can be called as a function with zero arguments, to
    set parameters to initial values
    """
    def __init__(s,pList):
        s.pList = pList
        s.__call__()
        return
    def __call__(s):
        for p in s.pList:
            p.init(0)
            pass
        return
    pass

class FinalParams:
    """constructor takes a list of ramped parameters. The constructor invokes
    each parameter such that it set to its final value.
    Also, this object can be called as a function with zero arguments, to
    set parameters to final values.
    """
    def __init__(s,pList):
        s.pList = pList
        s.__call__()
        return
    def __call__(s):
        for p in s.pList:
            p.finalize()
            pass
        return
    pass
        
from potList import PotList
class PotListWithContext(PotList):
    """a <m potList>.PotList with an initializing function called before
    calcEnergy/calcEnergyAndDerivs
    """
    def __init__(s,name="",
                 potList=None,
                 context=None):
        if not potList: potList = PotList()
        PotList.__init__(s,name)
        if potList: s.copy(potList)
        s.context = context
        from inspect import currentframe, getouterframes
        s.callingFrame = getouterframes( currentframe() )[1][0]
        return
    def calcEnergy(s):
        if s.context: s.context()
        return PotList.calcEnergy(s)
    def calcEnergyAndDerivs(s,derivs):
        if s.context: s.context()
        return PotList.calcEnergyAndDerivs(s,derivs)
    



class AnnealIVM:
    """class to perform simulated annealing using molecular dynamics. """
    def __init__(s,  initTemp,  finalTemp,
                 ivm=0,
                 numSteps=None,   tempStep=None,
                 rampedParams   ={},
                 extraCommands  =0,
                 toleranceFactor=1000):
        """construct by specifying the intial and final annealing temperatures,
        and an <m ivm>.IVM object.

        if tempStep is specified, it will be used to determine the number of
          dynamics runs at different temperatures. If if is omitted, numSteps
          will be used (and tempStep implicitly determined).

        rampedParams is a list of MultRamp and LinRamp objects which specify
        refinement parameters to adjust during the simulated annealing run.
        extraCommands is a function or string which is run before dynamics at
        each temperature. If it is a function, is is passed the current
        AnnealIVM instance as the argument.

        toleranceFactor is used to adjust the <m ivm>.IVM's energy tolerance
        as the temperature changes. The energy tolerance is calculated as
        
            eTolerance = temp / toleranceFactor

        The ivm argument doesn't actually have to be an <m ivm>.IVM, but it
        must be an object which has the following methods defined:
          setBathTemp
          setETolerance
          run

        """
        from inspect import currentframe, getouterframes
        s.initTemp = initTemp
        s.finalTemp = finalTemp
        s.extraCommands = extraCommands
        if tempStep:
            s.tempStep=tempStep
            s.numSteps = int( (initTemp - finalTemp)/float(tempStep) )
        elif numSteps:
            s.numSteps = numSteps
            s.tempStep = (initTemp - finalTemp)/float(numSteps)
        else:
            raise("AnnealIVM: neither numSteps nor tempStep is defined")
        s.ivm = ivm
        s.params = rampedParams
        s.toleranceFactor = toleranceFactor
        s.callingFrame = getouterframes( currentframe() )[1][0]
        return
    def run(s):
        s.bathTemp = s.initTemp
        s.printTemp()

        s.initParameters()
        
        s.runExtraCommands()
        s.runIVM(s.bathTemp) #run at initial temperature
        
        for i_cool in range(0,s.numSteps):
            s.bathTemp -= s.tempStep
            s.printTemp()

            s.updateParameters()
            s.runExtraCommands()

            s.runIVM(s.bathTemp)

            pass
        return

    def printTemp(s):
        import simulationWorld
        simWorld = simulationWorld.SimulationWorld_world()
        if simWorld.logLevel() != 'none':
            print "AnnealLoop: current temperature: %.2f" % s.bathTemp
            pass
        return

    def runExtraCommands(s):
        if s.extraCommands:
            if type(s.extraCommands) == type("string"):
                global_dict = s.callingFrame.f_globals
                local_dict = s.callingFrame.f_locals
                local_dict["annealLoopInfo"] = s
                exec( s.extraCommands, global_dict, local_dict )
                del local_dict["annealLoopInfo"]
            else:
                s.extraCommands(s)
                pass
            pass
        return
            
    def initParameters(s):
        "initialize ramped parameters"
        for param in s.params:
            param.init( s.numSteps, s ) #calls runAction
            pass
    def finalParameters(s):
        "sets parameters to final values"
        for param in s.params:
            param.val = param.stopValue
            param.runAction(s)
            pass
    def updateParameters(s):
        for param in s.params:
            param.update(s)
            pass
        return

    def runIVM(s,temp):

        if not s.ivm:
            return

        s.ivm.setBathTemp( temp )
        if s.toleranceFactor>0:
            s.ivm.setETolerance( temp / float(s.toleranceFactor) )
            pass
        s.ivm.run()
        
        return
    pass



def verticalFormat(name,
                   nameValuePairs,
                   minWidth=6,
                   numericFormat=".2f"):
    line2 = " " + name + ":"
    line1 = " " * len(line2)
    for (name,value) in nameValuePairs:
        strLen = max(minWidth,len(name))
        format = " %" + "%ds" % strLen
        line1 += format % name
        format = " %" + ("%d" % strLen) + numericFormat
        line2 += format % value
        pass
    return (line1, line2)

    
                   


if __name__ == "__main__":
    #
    # tests
    #

    
    from ensembleSimulation import EnsembleSimulation
    import sys
    import simulationWorld
    esim = EnsembleSimulation("esim",2)

    initSeed=751
    simWorld = simulationWorld.SimulationWorld_world()
    simWorld.setRandomSeed(751)

    result = ""
    expectedResult="""running structure 0 in process 0  thread: 0 seed: 751
running structure 1 in process 0  thread: 0 seed: 752
running structure 2 in process 0  thread: 0 seed: 753
running structure 3 in process 0  thread: 0 seed: 754
running structure 4 in process 0  thread: 0 seed: 755
running structure 5 in process 0  thread: 0 seed: 756
running structure 6 in process 0  thread: 0 seed: 757
running structure 7 in process 0  thread: 0 seed: 758
running structure 8 in process 0  thread: 0 seed: 759
running structure 9 in process 0  thread: 0 seed: 760
"""
    def concat(str):
        global result
        result += str
        return 0

    
    sys.stdout.write("StructureLoop: action as function...")
    StructureLoop(numStructures=10,
                  structLoopAction=lambda loopInfo: \
                  concat("running structure %d" % loopInfo.count +
                         " in process %d " % loopInfo.processID +
                         " thread: %d" % esim.member().memberIndex() +
                         " seed: %d\n" % loopInfo.randomSeed)
                  ).run()

    if result == expectedResult:
        print "ok"
    else:
        print "FAILED"
        pass


    simWorld.setRandomSeed(751)
    result=""
        
    sys.stdout.write("StructureLoop: action as string...")
        
    StructureLoop(numStructures=10,
                  structLoopAction=r'''global result
result += "running structure %d" % structLoopInfo.count + \
" in process %d " % structLoopInfo.processID + \
" thread: %d" % esim.member().memberIndex() + \
" seed: %d\n" % structLoopInfo.randomSeed
'''
                  ).run()
        
    if result == expectedResult:
        print "ok"
    else:
        print "FAILED"
        pass
            
    del esim

    sys.stdout.write("MultRamp: action as function...")
    expectedResult = "0.10 0.16 0.25 0.40 0.63 1.00 1.58 2.51 3.98 6.31 10.00 "
    result = ""
    def multProc(v):
        global result
        result +=  "%.2f " % v
        return
    
    param = MultRamp(0.1,10,multProc)
    numSteps=10
    param.init(numSteps)
    for i in range(numSteps):
        param.update()
        pass

    if result == expectedResult:
        print "ok"
    else:
        print "FAILED"
        pass
    

    sys.stdout.write("MultRamp: action as string...")
    expectedResult = "0.10 0.16 0.25 0.40 0.63 1.00 1.58 2.51 3.98 6.31 10.00 "
    result = ""
    def multProc(v):
        global result
        result +=  "%.2f " % v
        return
    
    param = MultRamp(0.1,10,"global result; result += '%.2f ' % VALUE")
    numSteps=10
    param.init(numSteps)
    for i in range(numSteps):
        param.update()
        pass

    if result == expectedResult:
        print "ok"
    else:
        print "FAILED"
        pass

    expectedResult="""0.10 extra: 1000.000000
0.32 extra: 900.000000
1.00 extra: 800.000000
3.16 extra: 700.000000
10.00 extra: 600.000000
"""

    sys.stdout.write("AnnealIVM: extraCmd as function...")
    result=""
        
    def coolProc(s):
        global result
        result += 'extra: %f\n' % s.bathTemp
        return

    AnnealIVM(initTemp=1000,
              finalTemp=600,
              tempStep=100,
              rampedParams=(param,),
              extraCommands=coolProc).run()

    if result == expectedResult:
        print "ok"
    else:
        print "FAILED"
        pass


    sys.stdout.write("AnnealIVM: extraCmd as string...")
    result=""
        
    AnnealIVM(initTemp=1000,
              finalTemp=600,
              tempStep=100,
              rampedParams=(param,),
              extraCommands=
              r"result += 'extra: %f\n' % annealLoopInfo.bathTemp").run()

    if result == expectedResult:
        print "ok"
    else:
        print "FAILED"
        pass

    pass

def testGradient(pots,
                 eachTerm=0,
                 sim=0):
    """
    check the analytic gradient in the given potential terms against finite
    difference values.

    If the eachTerm argument is set, each term in the potList is tested
    individually.
    """

    from potList import PotList
    potList = PotList()

    try:
        len(pots)
        for term in pots:
            potList.append(term)
            pass
        pass
    except:
        potList.append(pots)
        pass


    if not sim:
        from simulation import currentSimulation
        sim = currentSimulation()
        pass
    
    #is this an EnsembleSimulation?
    from ensembleSimulation import EnsembleSimulation_currentSimulation
    esim = EnsembleSimulation_currentSimulation()
    if esim and sim.name() == esim.name():
        sim = esim
    else:
        esim=0
        pass
    

    if eachTerm:
        ret=1
        for term in potList:
            print "testing gradient of potential term:", term.instanceName()
            if not testGradient(term):
                ret=0
            pass
        return ret
    
    from vec3 import Vec3

    # total gradient
    
    from derivList import DerivList
    dlist = DerivList()
    dlist.init(sim)
    energy = potList.calcEnergyAndDerivs(dlist)
    derivs = dlist.get(sim)
    
    eps=1e-7
    #eps=1e-8
    tol=1e-3

    ret=1
    if esim:
        from ensembleSharedObj import SharedObj
        sharedObj = SharedObj()

        for j in range( esim.size() ):
            member = esim.members(j)
            for i in range( esim.numAtoms() ):
                if j == esim.member().memberIndex():
                    initPos = Vec3( member.atomPos(i) )
                    pos = Vec3( initPos )
                    pos[0] += eps
                    member.setAtomPos(i,pos)
                    pass
                sharedObj.set( 0 )
                denergy = potList.calcEnergy().energy-energy.energy
                if j == esim.member().memberIndex():
                    grad = denergy/eps
                    if abs(grad-derivs[i][0])>tol*max(abs(grad),
                                                      abs(derivs[i][0])):
                        sharedObj.set( (grad, derivs[i][0]) )
                        pass
                    esim.setAtomPos(i,initPos)
                    pass
                if sharedObj.barrierGet():
                    ret=0
                    print i,j,sharedObj.barrierGet()
                    pass
                pass
            pass
        pass
    else:
        #not an EsembleSimulation
        for i in range( sim.numAtoms() ):
            initPos = Vec3( member.atomPos(i) )
            pos = Vec3( initPos )
            pos[0] += eps
            member.setAtomPos(i,pos)
            denergy = potList.calcEnergy().energy-energy.energy
            grad = denergy/eps
            if abs(grad-derivs[i][0])>tol:
                ret=0
                print i,j,grad, derivs[i][0]
                pass
            sim.setAtomPos(i,initPos)
            pass
        pass
    return ret

    

#default threshold values
xplorThreshold = {}
xplorThreshold['BOND'] = 0.05
xplorThreshold['ANGL'] = 2
xplorThreshold['IMPR'] = 2
xplorThreshold['CDIH'] = 5
xplorThreshold['NOE']  = 0.5
xplorThreshold['COUP'] = 1.
xplorThreshold['HBDA'] = 0.
xplorThreshold['VDW'] = 0.2

from cdsMatrix import CDSMatrix_double
chemTypeDist = CDSMatrix_double(0,0,0.)
from cdsVector import CDSVector_int
lookupVector = CDSVector_int(0)
bondList=[]

def initializeRadii():
    from tempfile import mktemp #FIX: possible security problem
    import xplor, os
    tmpFilename=mktemp(suffix=".xplor")+fileSuffix
    xplor.command("set echo off mess off end")
    xplor.command("write params output=%s end" % tmpFilename)
    xplor.command("set echo $prev_echo mess $prev_messages end")
    readXplorRadii(tmpFilename)
    os.unlink(tmpFilename)
    return
def readXplorRadii(filename):
    import re
    global chemTypeDist
    chemTypeRadius=[]
    chemTypeLookup = {}
    for line in open(filename).readlines():
        match = re.search(r"^\s*nonb.*\s+([a-z0-9]+)\s+([0-9.]+)\s+([0-9.]+)",
                          line,re.IGNORECASE)
        if match:
            chemType = match.group(1)
            sigma = float( match.group(3) )
            i = len(chemTypeRadius)
            chemTypeRadius.append( sigma )
            chemTypeLookup[chemType] = i
            chemTypeDist.resize(i+1,i+1)
            for j in range(len(chemTypeRadius)):
                chemTypeDist[i,j] = 0.5**(5./6) * (chemTypeRadius[i]+
                                                   chemTypeRadius[j])
                chemTypeDist[j,i] = chemTypeDist[i,j]
                pass
            pass

        match = re.search(
            r"^\s*nbfi.*\s+([a-z0-9]+)\s+([a-z0-9]+)\s+([0-9.]+)\s+([0-9.]+)",
                          line,re.IGNORECASE)
        if match:
            i = chemTypeLookup[ match.group(1) ]
            j = chemTypeLookup[ match.group(2) ]
            A = float( match.group(3) )
            B = float( match.group(4) )
            sigma = (A/B)**(1./6)
            chemTypeDist[i,j] = 0.5**(5./6) * sigma
            chemTypeDist[j,i] = chemTypeDist[i,j]
            pass
        pass

    from xplor import simulation as sim
    lookupVector.resize( sim.numAtoms() )
    for i in range(sim.numAtoms()):
        lookupVector[i] = chemTypeLookup[ sim.atomByID(i).chemType() ]
        pass

    global bondList
    bondList=[]
    for i in range(sim.numAtoms()):
        bondList.append([])
        pass

    import xplor
    xplor.command("set echo off mess off end")
    nbondNBXMod = int(xplor.command("param nbond ? end end","NBXMOD")[0])
    xplor.command("set echo $prev_echo mess $prev_messages end")
    if abs(nbondNBXMod)>1:
        for i in range(sim.numBonds()):
            pair = sim.bondPairByID(i)
            bondList[pair[0]].append( pair[1] )
            bondList[pair[1]].append( pair[0] )
            pass
        pass
    #add in atoms with 1-3 relationship
    if abs(nbondNBXMod)>2:
        newList=[]
        for i in range(sim.numAtoms()):
            newList.append(bondList[i][:])
            pass
        for i in range(sim.numAtoms()):
            for j in bondList[i]:
                for k in bondList[j]:
                    if k != i and not k in newList:
                        newList[i].append(k)
                        pass
                    pass
                pass
            pass
        bondList = newList
        pass
    #add in atoms with 1-4 relationship
    if abs(nbondNBXMod)>3:
        newList=[]
        for i in range(sim.numAtoms()):
            newList.append(bondList[i][:])
            pass
        for i in range(sim.numAtoms()):
            for j in bondList[i]:
                for k in bondList[j]:
                    if k != i and not k in newList:
                        newList[i].append(k)
                        pass
                    pass
                pass
            pass
        bondList = newList
        pass
    return

def vdwViolations(threshold):
    """determine nbonded violations:
      atoms closer than the appropriate vdw radius.
    """
    from xplor import simulation as sim
    if len(lookupVector) != sim.numAtoms():
        initializeRadii()
        

    from atomSel import AtomSel
    atoms = AtomSel("not resname ANI")

    import xplor
    xplor.command("set echo off mess off end")
    radiusScale = float(xplor.command("param nbond ? end end","repel")[0])
    xplor.command("set echo $prev_echo mess $prev_messages end")
    distanceMatrix = CDSMatrix_double(chemTypeDist)
    distanceMatrix.scale( radiusScale )

    print
    print "  Nonbonded Violations"
    print "        threshold: %f" % threshold
    print
    print "%26s %26s %6s %6s" % ("atom1     ", "atom2     ",
                                 "dist", "vdwDist")
    print "_"*70

    violations=[]
    from vec3 import norm
    from nbond import findCloseAtoms
    for (i,j) in findCloseAtoms(atoms,lookupVector,
                                distanceMatrix, threshold):
        if not j in bondList[i]:
            atomi = atoms[i]
            atomj = atoms[j]
            violations.append( (atomi, atomj,
                                norm(atomi.pos()-atomj.pos()),
                                distanceMatrix[lookupVector[i],
                                               lookupVector[j]]) )
            pass
        pass

    for (a1,a2,dist,vdwDist) in violations:
        print "%26s %26s %6.2f %6.2f" % (a1.string(), a2.string(),
                                         dist, vdwDist)
        pass
        
    return len(violations)

def xplorPot_summarize(potList,outFilename=0):
    """get rms, violations summaries from XplorPots.
    Print to outFilename if it is set."""
    potList = filter(lambda x: x.potName().find('XplorPot')>=0,potList)
    terms = []
    import xplor, sys

    xplor.command("set echo off mess off end")

    import tempfile, os
        
    if outFilename:
        old_stdout = sys.stdout
        sys.stdout = open(outFilename,"a")
        pass

    #standard terms
    for name in ('BOND','ANGL','IMPR','CDIH','NOE'):
        if filter(lambda x: x.instanceName()==name,potList):
            print "\n   Violations in XPLOR term %8s" % name
            print "             threshold: %f\n" % xplorThreshold[name]
            if outFilename:
                xplorFilename=tempfile.mktemp('xplor-violations.' +
                                              fileSuffix)
                xplor.command("set print %s end" % xplorFilename)
                pass
            (rms,violations) = xplor.command("print thres=%f %s" %
                                             (xplorThreshold[name],name),
                                             ("result","violations"))
            if outFilename:
                xplor.command("set print OUTPUT end" )
                xplor.command("close %s end" % xplorFilename)
                print open(xplorFilename).read()
                os.unlink(xplorFilename)
                pass
            terms.append( ('Xplor',name,float(rms),int(violations)) )
        pass

    #inverted form
    for name in ('COUP','HBDA'):
        if filter(lambda x: x.instanceName()==name,potList):
            print "\n   Violations in XPLOR term %8s" % name
            print "             threshold: %f\n" % xplorThreshold[name]
            if outFilename:
                xplorFilename=tempfile.mktemp('xplor-violations.'+
                                              fileSuffix)
                xplor.command("set print %s end" % xplorFilename)
                pass
            (rms,violations) = xplor.command("%s print thres=%f end end" %
                                             (name,xplorThreshold[name]),
                                             ("result","violations"))
            if outFilename:
                xplor.command("set print OUTPUT end" )
                xplor.command("close %s end" % xplorFilename)
                print open(xplorFilename).read()
                os.unlink(xplorFilename)
                pass
            terms.append( ('Xplor',name,float(rms),int(violations)) )
        pass

    xplor.command("set echo $prev_echo mess $prev_messages end")

    #vdw
    if filter(lambda x: x.instanceName()=='VDW',potList):
        viols = vdwViolations(xplorThreshold['VDW'])
        terms.append( ('Xplor','VDW',-1,viols) )
        pass

    if outFilename:
        sys.stdout.close()
        sys.stdout = old_stdout
        pass

    #all others
    for pot in potList:
        name = pot.instanceName()
        if not name in map(lambda x:x[1],terms):
            terms.append( ('Xplor',name,-1,-1) )
            
    return terms

def nativePot_summarize(potList):
    "get rms, violations summaries from non-XplorPots"
    potList = filter(lambda x: x.potName().find('XplorPot')<0,potList)
    terms = []

    for pot in potList:
        rms=-1
        viols=-1
        if pot.potName()=='PotList' and len(pot):
            lterms=nativePot_summarize(pot)
            rms=0.
            viols=0
            count=0
            for (pName,name,lrms,lviols) in lterms:
                if lrms>=0:
                    count += 1
                    rms += lrms
                    viols += lviols
                    pass
                pass
            if count>0:
                rms /= count
                viols
                pass
        else:
            try:
                rms = pot.rms()
            except AttributeError:
                pass
            try:
                viols = pot.violations()
            except AttributeError:
                pass
            pass
        terms.append((pot.potName(),pot.instanceName(),rms,viols))
        pass
    return terms

from rdcPotTools import RDC_analyze
from csaPotTools import CSA_analyze
from varTensorTools import VarTensor_analyze

def flattenPotList(potList):
    from potList import PotListPtr
    ret = []
    for pot in potList:
        if pot.potName()=='PotList':
            p=PotListPtr( pot.this )
            ret += flattenPotList(p)
        else:
            ret.append(pot)
            pass
        pass
    return ret


def analyze(potList,extraTerms=PotList(),
            outFilename=0):
    """ pretty print appropriate terms from the given PotList and return
    them in a remarks string.
    The optional extraTerms is a PotList of terms which do not
    contribute to the total energy, but which should be analyzed all the same-
    use this for cross-validated potential terms.

    If outFilename is specified, the verbose violations info is written to
    the specified file.

    """

    if outFilename:
        outfile = open(outFilename,"w")
        print >> outfile, "\n  Violation Analysis \n" 
        outfile.close()
        pass

    ret  = "-"*60 + '\n'
    ret += "summary    Name       Energy      RMS     Violations\n"
    ret += "summary %-10s %10.2f\n" % ("total",
                                      potList.calcEnergy().energy)

    reports = potList.energyReports()

    terms = xplorPot_summarize(potList,outFilename)
    terms += nativePot_summarize(potList)
    
    #terms += NOEPot_summarize(potList)
    
    for term in terms:
        (type,name,rms,viols) = term
        energy = filter(lambda x: x[0]==name,reports)[0][1]
        rmsString=" "*8
        if rms>-1: rmsString = "%8.3f" % rms
        violString=' '*8
        if viols>-1: violString = "%8d" % viols
        ret += "summary %-10s %10.2f %s %s\n" % (term[1],
                                               energy,rmsString,violString)
        pass

    ret += "-"*60 + '\n'
    if len(extraTerms):
        ret += "\nCross-validated terms:\n"
        extraTerms.calcEnergy()
        reports = extraTerms.energyReports()

        terms = xplorPot_summarize(extraTerms,outFilename)
        terms += nativePot_summarize(extraTerms)
    
    
        for term in terms:
            (type,name,rms,viols) = term
            energy = filter(lambda x: x[0]==name,reports)[0][1]
            rmsString=" "*8
            if rms>-1: rmsString = "%8.3f" % rms
            violString=' '*8
            if viols>-1: violString = "%8d" % viols
            ret += "summary %-10s %10.2f %s %s\n" % (term[1],
                                                     energy,
                                                     rmsString,violString)
            pass

        terms = flattenPotList(extraTerms)

        if len(terms) > len(extraTerms):
            ret += "summary\n"
        
            lineBeg = "summary cross-validated terms: "
            line=lineBeg
            for term in terms:
                if (len(line)+len(term.instanceName()))>71: 
                    ret += line + '\n'
                    line=lineBeg
                    pass
                line += "%s" % term.instanceName()
                if term != terms[-1]: line += ", "
                pass
            if line != lineBeg: ret += line + '\n'
            pass

        ret += "-"*60 + '\n'

    if outFilename:
        import sys
        old_stdout = sys.stdout
        sys.stdout = open(outFilename,"a")
        pass
        

    import noePotTools
    tmp = noePotTools.analyze(list(potList)+list(extraTerms))
    if tmp:
        ret += ' NOE terms\n'
        ret += reduce(lambda x,y: x+'\n'+y,map(lambda x: "NOE "+x,
                                               tmp.splitlines()))
        ret += '\n' + "-"*60 + '\n'
        pass

    tmp = RDC_analyze( list(potList)+list(extraTerms) )
    if tmp:
        ret += " Dipolar Coupling Analysis\n"
        ret += reduce(lambda x,y: x+'\n'+y,map(lambda x: "RDC "+x,
                                               tmp.splitlines()))
        ret += '\n' + "-"*60 + '\n'
        pass

    tmp = CSA_analyze( list(potList)+list(extraTerms) )
    if tmp:
        ret += " CSA Analysis\n"
        ret += reduce(lambda x,y: x+'\n'+y,map(lambda x: "CSA "+x,
                                               tmp.splitlines()))
        ret += '\n' + "-"*60 + '\n'
        pass

    tmp = VarTensor_analyze(list(potList)+list(extraTerms))
    if tmp:
        ret += ' Variable Tensor Analysis\n'
        ret += reduce(lambda x,y: x+'\n'+y,map(lambda x: "VarTensor "+x,
                                               tmp.splitlines()))
        ret += '\n' + "-"*60 + '\n'
        pass

    import jCoupPotTools
    tmp = jCoupPotTools.analyze(list(potList)+list(extraTerms))
    if tmp:
        ret += ' J-Coupling terms\n'
        ret += reduce(lambda x,y: x+'\n'+y,map(lambda x: "JCoup "+x,
                                               tmp.splitlines()))
        ret += '\n' + "-"*60 + '\n'
        pass

    import shapePotTools
    tmp = shapePotTools.analyze(list(potList)+list(extraTerms))
    if tmp:
        ret += ' Shape Tensor Analysis\n'
        ret += reduce(lambda x,y: x+'\n'+y,map(lambda x: "Shape "+x,
                                               tmp.splitlines()))
        ret += '\n' + "-"*60 + '\n'
        pass

    import posRMSDPotTools
    tmp = posRMSDPotTools.analyze(list(potList)+list(extraTerms))
    if tmp:
        ret += ' Positional RMSD Analysis\n'
        ret += reduce(lambda x,y: x+'\n'+y,map(lambda x: "PosRMSD "+x,
                                               tmp.splitlines()))
        ret += '\n' + "-"*60 + '\n'
        pass

    
    if outFilename:
        import sys
        sys.stdout.close()
        sys.stdout = old_stdout
        pass
    
    # first print overall energies, energies of each term

    #go through each potential type for analysis:
    # print out results if verbose flag is set.
    # bond
    # angle
    # impropers
    #

    # rms, # violations

    # NOE

    # rdc:
    #
    # deal with ensemble, non-ensemble cases
    return ret