import sys
xplor.requireVersion("2.17")

numberOfStructures=1
startStructure=0
ensembleSize=8          # number of ensemble members
outFilename = "OUTPUT/%d_Random_STRUCTURE_MEMBER.pdb" % ensembleSize

import protocol
protocol.initRandomSeed()

#
# read in the PSF and initial PDB files
#
protocol.initStruct("1ubq_template.psf")
protocol.initCoords("1ubq_template.pdb")
protocol.initParams("./TopPar/parallhdg_new.pro")
protocol.initParams("./TopPar/parnah1er1_mod_new.inp")

from ensembleSimulation import EnsembleSimulation
esim = EnsembleSimulation("ensemble",ensembleSize)

# list of potential terms used in refinement
from potList import PotList
potList = PotList()

# parameters to ramp up during the simulated annealing protocol
#
from simulationTools import MultRamp, StaticRamp, InitialParams
rampedParams=[]

protocol.covalentMinimize("not resname ANI")

from avePot import AvePot
from xplorPot import XplorPot

from rdcPotTools import Da_prefactor

protocol.initNBond(cutnb=4.5)
potList.append( AvePot(XplorPot,"VDW") )
rampedParams.append( MultRamp(0.9,0.78,
                              "xplor.command('param nbonds repel VALUE end end')") )
rampedParams.append( MultRamp(0.004,4,
                              "xplor.command('param nbonds rcon VALUE end end')") )

for name in ("bond","angl","impr"):
    potList.append( AvePot(XplorPot,name) )
    pass
rampedParams.append( MultRamp(0.4,1.0,"potList['ANGL'].setScale(VALUE)"))
rampedParams.append( MultRamp(0.1,1.0,"potList['IMPR'].setScale(VALUE)"))


from ivm import IVM
import varTensorTools
mini = IVM()             #initial alignment of orientation tensor axes
#map(lambda x: mini.potList().append(x), (rdcPots, csaPots) )

protocol.initMinimize(mini,
                      numSteps=20)
mini.fix("not resname ANI")
#mini.run()               #this initial minimization is not strictly necessary

dyn = IVM()
protocol.initDynamics(dyn,potList=potList)
protocol.torsionTopology(dyn)

# Give atoms uniform weights, except for the anisotropy axis
from atomAction import SetProperty
AtomSel("not resname ANI").apply( SetProperty("mass",100.) )
import varTensorTools
AtomSel("all            ").apply( SetProperty("fric",10.) )

##
## minc used for final cartesian minimization
##
from selectTools import IVM_groupRigidSidechain
minc = IVM()
protocol.initMinimize(minc,potList=potList)
IVM_groupRigidSidechain(minc)
protocol.cartesianTopology(minc,"not resname ANI")

init_t = 3000
annealTime=1.5     #time to integrate at a given annealing temp.
annealSteps=0    #max num steps to be taken during a single annealing temp.

class DynMin:
    """class to perform minimization followed by dynamics
    """
    def __init__(s,ivm):
        s.ivm = ivm
        return
    def setBathTemp(s,temp):
        s.ivm.setBathTemp(temp)
        return
    def setETolerance(s,tol):
        s.ivm.setETolerance(tol)
        return
    def run(s):
        #protocol.initMinimize(s.ivm,
        #                      numSteps=1000)
        #s.ivm.run()
        protocol.initDynamics(s.ivm,
                              finalTime=annealTime,
                              numSteps=annealSteps,
                              eTol_minimum=0.001
                              )
        s.ivm.run()
        return

from simulationTools import AnnealIVM
anneal= AnnealIVM(initTemp =init_t,
                  finalTemp=25,
                  tempStep =25,
                  ivm=DynMin(dyn),
                  rampedParams = rampedParams)
    

# initialize parameters for initial minimization.
InitialParams( rampedParams )

# initial minimization
protocol.initMinimize(dyn,
                      numSteps=1000)
dyn.run()

#short symmetry-breaking dynamics run
protocol.initDynamics(dyn,
                      bathTemp=init_t,
                      initVelocities=1,
                      finalTime=0.1,
                      numSteps=1000)
dyn.run()
# minimize symmetry-broken ensemble
protocol.initMinimize(dyn,
                      numSteps=1000)
dyn.run()

from simulationTools import testGradient
#testGradient(potList,eachTerm=1)

from pdbTool import PDBTool

def calcOneStructure(loopInfo):
    # initialize parameters for high temp dynamics.
    InitialParams( rampedParams )

    # high temperature bit - using only P-P nonbonded terms
    protocol.initNBond(repel=1.2,
                       cutnb=100,
                       tolerance=45,
                       selStr="name P")
    protocol.initDynamics(dyn,
                          initVelocities=1,
                          bathTemp=init_t,
                          potList=potList,
                          numSteps=5000,
                          finalTime=10)
    dyn.run()
    protocol.initNBond() #reset to include all atoms

    # perform simulated annealing
    #
    anneal.run()
              
    #
    # torsion angle minimization
    #
    protocol.initMinimize(dyn)
    dyn.run()

    ##
    ##all atom minimization
    ##
    minc.run()

    #
    # perform analysis and write structure
    loopInfo.writeStructure(potList)

    #calculate regularized ensemble ave. struct.
    aveFilename = outFilename.replace('MEMBER','ave')

    savedCoords = esim.atomPosArr()
    esim.setAtomPosArr( esim.meanAtomPosArr() )

    from simulationTools import minimizeRefine
    minimizeRefine(potList=potList,
                   refineSteps=50)

    if esim.singleThread():
        PDBTool(loopInfo.makeFilename(aveFilename),
                "not resname ANI").write()
        
        pass
    esim.multiThread()
    esim.setAtomPosArr( savedCoords )

    pass


from simulationTools import StructureLoop, FinalParams
StructureLoop(numStructures=numberOfStructures,
              startStructure=startStructure,
              structLoopAction=calcOneStructure,
              pdbTemplate=outFilename,
              genViolationStats=1).run()
              #averageTopFraction=0.5,
              #averagePotList=potList,
              #averageCrossTerms=crossTerms,
              #averageContext=FinalParams(rampedParams),
              #averageFilename="ave.pdb",
              #averageFitSel="not name H* and not resname ANI",
              #averageCompSel="not name H* and not resname ANI"       ).run()