#
# slow cooling protocol in torsion angle space for protein G. Uses 
# NOE, RDC, J-coupling restraints.
#
# CDS 5/14/03
# modified 7/2/03
#

# this checks for typos on the command-line. User-customized arguments could
# also be specified.
#
xplor.parseArguments()

outFilename = "SCRIPT_STRUCTURE.sa"
numberOfStructures=2
seed=3421
simWorld.setRandomSeed(seed)   #set random seed

command = xplor.command

#
# import necessary modules
#
from xplorPot import XplorPot
from rdcPotTools import create_RDCPot
from varTensorTools import create_VarTensor
import varTensorTools
from ivm import IVM
from potList import PotList
import protocol

#
# generate PSF data from sequence and initialize the correct parameters.
#
from psfGen import seqToPSF
seqToPSF(open('protG.seq').read())

#
# generate a random extended structure with correct covalent geometry
#
protocol.genExtendedStructure("gb1_extended_%d.pdb" % seed)

#
# a PotList conatins a list of potential terms. This is used to specify which
# terms are active during refinement.
#
potList = PotList()

# orientation Tensor - used with the dipolar coupling term
#
oTensor = create_VarTensor("oTensor")
oTensor.setDa(-9.9)
oTensor.setRh(0.23)

# dipolar coupling restraints for protein amide NH.  
#
rdc = create_RDCPot(name="rdc",
                    oTensor=oTensor,
                    file="bicelles_new_nh.tbl")
rdc.setThreshold(0.5)
print rdc.info()
potList.append(rdc)


# set up NOE potential
from noePotTools import create_NOEPot
noe = create_NOEPot("noe",file="protG_NOE.tbl")
noe.setPotType( "soft" )
noe.setRSwitch( 0.5 )
noe.setAsympSlope( 1. )
noe.setSoftExp(1.)
noe.setThreshold(0.5)
print noe.info()
potList.append(noe)

# set up J coupling
from jCoupPotTools import create_JCoupPot
jCoup = create_JCoupPot("jcoup",file="jna_coup.tbl",
                        A=6.98,B=-1.38,C=1.72,phase=-60)
jCoup.setThreshold(1.0)
print jCoup.info()
potList.append(jCoup)

# Set up dihedral angles
protocol.initDihedrals("protG_dih.tbl",
                       scale=5,          #initial force constant
                       useDefaults=0)
potList.append( XplorPot('CDIH') )

# radius of gyration term - not used here.
#
#protocol.initCollapse("resid 4:20",
#                      scale=25.0,
#                      Rtarget=14.0)
#potList.append( XplorPot('COLL') )



#
# setup parameters for atom-atom repulsive term. (van der Waals-like term)
#
protocol.initNBond(nbxmod=4,  # Can use 4 here, due to IVM dynamic
                   repel=0.5) # initial effective atom radius
potList.append( XplorPot('VDW') )

#
# annealing settings
#

init_t  = 3500.     # Need high temp and slow annealing to converge
final_t = 100

cool_steps = 15000   # slow annealing


#
# initial force constant settings
#
ini_rad  = 0.4 
ini_con = 0.004
ini_ang = 0.4
ini_imp = 0.4     # was 0.1
ini_noe = 1       # was 150
ini_sani = 0.01
     
    
potList.add( XplorPot("BOND") )
potList.add( XplorPot("ANGL") )
potList.add( XplorPot("IMPR") )
      

#
# potential terms used for high-temp dynamics
#
hitemp_potList = PotList()
hitemp_potList.add( XplorPot("BOND") )
hitemp_potList.add( XplorPot("ANGL") )
hitemp_potList.add( XplorPot("IMPR") )
hitemp_potList.add( XplorPot("CDIH") )
hitemp_potList.add( noe   )
hitemp_potList.add( jCoup )
hitemp_potList.add( rdc   )

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

#
# IVM setup
#

dyn = IVM()

# Minimize in Cartesian space with only the covalent constraints.
# Note that bonds, angles and many impropers can't change with the 
# internal torsion-angle dynamics
#  breaks bonds topologically - doesn't change force field

dyn.potList().add( XplorPot("BOND") )
dyn.potList().add( XplorPot("ANGL") )
dyn.potList().add( XplorPot("IMPR") )

dyn.breakAllBondsIn("not resname ANI")
oTensor.setFreedom("fix")
varTensorTools.topologySetup(dyn,oTensor)

protocol.initMinimize(dyn,numSteps=1000)
dyn.run()

#
# reset ivm topology for torsion-angle dynamics
#
dyn.reset()
dyn.potList().removeAll()

oTensor.setFreedom("fixDa, fixRh")
varTensorTools.topologySetup(dyn,oTensor)
protocol.torsionTopology(dyn)

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

oTensor.setFreedom("varyDa, varyRh")
varTensorTools.topologySetup(minc,oTensor)


def setConstraints(k_ang,k_imp):
    command("""
      constraints 
         interaction (not resname ANI) (not resname ANI) 
            weights * 1 angl %f impr %f
         end 
      end""" % ( k_ang, k_imp) )
    return


def structLoopAction(loopInfo):
    #
    # this function calculates a single structure.
    #

    #
    # set some high-temp force constants
    #
    noe.setScale( 20 )       #use large scale factor initially
    rdc.setScale( ini_sani ) 
    command("parameters  nbonds repel %f end end" % ini_rad)
    command("parameters  nbonds rcon %f end end" % ini_con)
    setConstraints(ini_ang, ini_imp)
    # Initial weight--modified later
    command("restraints dihedral scale=5. end")

    #
    # high temp dynamics
    # note no Van der Waals term
    #

    init_t = 3500
    ini_timestep = 0.010
    bath = init_t
    timestep = ini_timestep

    protocol.initDynamics(dyn,
                          potList=hitemp_potList, # potential terms to use
                          bathTemp=bath,
                          initVelocities=1,
                          finalTime=20,
                          printInterval=100)

    dyn.setETolerance( bath/100 )  #used to det. stepsize. default: bath/1000 
    dyn.run()

    # increase dihedral term 
    command("restraints dihedral scale=200. end")

    #
    # cooling and ramping parameters
    #

    global k_ang, k_imp

    # MultRamp ramps the scale factors over the specified range for the 
    # simulated annealing.
    from simulationTools import MultRamp
    k_noe = MultRamp(ini_noe,30.,"noe.setScale( VALUE )")
    k_rdc = MultRamp(ini_sani,1.,"rdc.setScale( VALUE )")
    radius = MultRamp(ini_rad,0.8,
                      "command('param nbonds repel VALUE end end')")
    k_vdw  = MultRamp(ini_con,4, "command('param nbonds rcon VALUE end end')")
    k_ang  = MultRamp(ini_ang,1.0)
    k_imp  = MultRamp(ini_imp,1.0)
    
    protocol.initDynamics(dyn,
                          potList=potList,
                          bathTemp=bath,
                          initVelocities=1,
                          finalTime=2,
                          printInterval=100,
                          stepsize=timestep)

    from simulationTools import AnnealIVM
    AnnealIVM(initTemp =init_t,
              finalTemp=100,
              tempStep =25,
              ivm=dyn,
              rampedParams = [k_noe,k_rdc,radius,k_vdw,k_ang,k_imp],
              extraCommands= lambda notUsed: \
                setConstraints(k_ang.value(),
                               k_imp.value())).run()
              
              
    #
    # final torsion angle minimization
    #
    protocol.initMinimize(dyn,
                          printInterval=50)
    dyn.run()

    #
    # final all atom minimization
    #
    protocol.initMinimize(minc,
                          potList=potList,
                          printInterval=100)
    minc.run()

    #
    # analyze and write out structure
    # 
    loopInfo.writeStructure(potList)
    pass



from simulationTools import StructureLoop
StructureLoop(numStructures=numberOfStructures,
              pdbTemplate=outFilename,
              structLoopAction=structLoopAction,
              genViolationStats=1,
              averagePotList=potList).run()