xplor.requireVersion("2.18")

xplor.parseArguments() # check for typos on the command-line

outFilename = "SCRIPT_STRUCTURE.pdb"
numberOfStructures=10

# protocol module has many high-level helper functions.
#
import protocol
protocol.initRandomSeed(3421)   #explicitly set random seed

from xplor import select

#
# annealing settings
#

command = xplor.command

protocol.initParams(["protein", "cysp.par"])

# read an existing model
#
protocol.initStruct(["c2ab23r2.psf", "addAtoms24.psf", "addAtoms24ba.psf"])
protocol.initCoords(["2pln2loen.pdb", "addAtoms24b.pdb", "addAtoms24ba.pdb"])
xplor.command("delete selection=(not known) end")


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

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

rampedParams=[]
highTempParams=[]

# set up NOE potential
noe=PotList('noe')
potList.append(noe)
from noePotTools import create_NOEPot
for (name,scale,file) in [('all',1,"dist_rstrnts_mem3.tbl"),
                          #add entries for additional tables
                          ]:
    pot = create_NOEPot(name,file)
    # pot.setPotType("soft") - if you think there may be bad NOEs
    pot.setScale(scale)
    noe.append(pot)
rampedParams.append( MultRamp(5,30, "noe.setScale( VALUE )") )

# Set up PlaneDist Potential terms
planePot=PotList('planeDist')
import planeDistTools
from planeDistTools import create_PlaneDistPot
aa = "resid 500 and name CE1"
ab = "resid 500 and name CZ"
ac = "resid 500 and name HZ"
ba = "resid 501 and name CE1"
bb = "resid 501 and name CZ"
bc = "resid 501 and name HZ"
for (name,file,oAtom,xAtom,yAtom) in [('plane1','a.tbl',aa,ab,ac),
                                      ('plane2','b.tbl',ba,bb,bc)]:
     pot = create_PlaneDistPot(name,restraintsFile=file,
                               oAtom=oAtom,
                               xAtom=xAtom,
                               yAtom=yAtom)
     pot.setShowAllRestraints(True)
     pot.setUseSign(True)
     pot.setFreedom('fix')
     planePot.append(pot)
     pass
rampedParams.append( MultRamp(0.1,30, "planePot.setScale( VALUE )") )
potList.append(planePot)


from xplorPot import XplorPot

#Rama torsion angle database
#
protocol.initRamaDatabase()
potList.append( XplorPot('RAMA') )
rampedParams.append( MultRamp(.002,1,"potList['RAMA'].setScale(VALUE)") )

#
# setup parameters for atom-atom repulsive term. (van der Waals-like term)
#
potList.append( XplorPot('VDW') )
rampedParams.append( StaticRamp("protocol.initNBond()") )
rampedParams.append( MultRamp(0.9,0.8,
                              "command('param nbonds repel VALUE end end')") )
rampedParams.append( MultRamp(.004,4,
                              "command('param nbonds rcon VALUE end end')") )
# nonbonded interaction only between CA atoms
highTempParams.append( StaticRamp("""protocol.initNBond(cutnb=100,
                                                        tolerance=45,
                                                        repel=1.2,
                                                        onlyCA=1)""") )


potList.append( XplorPot("BOND") )
potList.append( XplorPot("ANGL") )
potList['ANGL'].setThreshold( 5 )
rampedParams.append( MultRamp(0.4,1,"potList['ANGL'].setScale(VALUE)") )
potList.append( XplorPot("IMPR") )
potList['IMPR'].setThreshold( 5 )
rampedParams.append( MultRamp(0.1,1,"potList['IMPR'].setScale(VALUE)") )
      


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


# IVM setup
#   the IVM is used for performing dynamics and minimization in torsion-angle
#   space, and in Cartesian space.
#
from ivm import IVM
dyn = IVM()
#
# minc used for final cartesian minimization
#
minc = IVM()

def accept(potList):
    """
    return True if current structure meets acceptance criteria
    """
    if potList['noe'].violations()>0:
        return False
#    if potList['rdc'].rms()>1.0: #this might be tightened some
#        return False
#    if potList['CDIH'].violations()>0:
#        return False
    if potList['BOND'].violations()>0:
        return False
    if potList['ANGL'].violations()>0:
        return False
    if potList['IMPR'].violations()>1:
        return False
    
    return True

def calcOneStructure(loopInfo):
    """ this function calculates a single structure, performs analysis on the
    structure, and then writes out a pdb file, with remarks.
    """

    InitialParams( rampedParams )
    dyn.reset()
    dyn.setGroupList([])         # set up a empty grouplist
    groupList = dyn.groupList()
    groupList.append( select('''resid 140:265 and (name ca or name n or 
                              name c or name o or name hn or name ha)'''))
    groupList.append( select('''(resid 171 or resid 173 or resid 174 or 
                              resid 189 or resid 199 or resid 202 or 
                              resid 213 or resid 231 or resid 234 or 
                              resid 236 or resid 239 or resid 244 or 
                              resid 300 or resid 304 or resid 305 or 
                              resid 306 or resid 327 or resid 332 or 
                              resid 334 or resid 367 or resid 368 or 
                              resid 369 or resid 370) and
                              (name cb or name hb1 or name hb2 or 
                              name sg1 or name sg2)'''))
    groupList.append( select('''resid 273:419 and (name ca or name n or
                              name c or name o or name hn or name ha)'''))
    groupList.append( select('''(resid 171 or resid 173 or resid 174 or
                              resid 189 or resid 199 or resid 202 or
                              resid 213 or resid 231 or resid 234 or
                              resid 236 or resid 239 or resid 244 or
                              resid 300 or resid 304 or resid 305 or
                              resid 306 or resid 327 or resid 332 or
                              resid 334 or resid 367 or resid 368 or
                              resid 369 or resid 370) and
                              (name cb or name hb1 or name hb2 or
                              name sg1 or name sg2)'''))
    dyn.setGroupList( groupList )
    planeDistTools.topologySetup(dyn,planePot)
    protocol.torsionTopology(dyn)
    # minc used for final cartesian minimization
    #
    #minc.fix( select('resid 200') )
    minc.setGroupList([])         # set up a empty grouplist
    groupList = minc.groupList()
    groupList.append( select('''resid 140:265 and (name ca or name n or
                              name c or name o or name hn or name ha)'''))
    groupList.append( select('''(resid 171 or resid 173 or resid 174 or
                              resid 189 or resid 199 or resid 202 or
                              resid 213 or resid 231 or resid 234 or
                              resid 236 or resid 239 or resid 244 or
                              resid 300 or resid 304 or resid 305 or
                              resid 306 or resid 327 or resid 332 or
                              resid 334 or resid 367 or resid 368 or
                              resid 369 or resid 370) and
                              (name cb or name hb1 or name hb2 or
                              name sg1 or name sg2)'''))
    groupList.append( select('''resid 273:419 and (name ca or name n or
                              name c or name o or name hn or name ha)'''))
    groupList.append( select('''(resid 171 or resid 173 or resid 174 or
                              resid 189 or resid 199 or resid 202 or
                              resid 213 or resid 231 or resid 234 or
                              resid 236 or resid 239 or resid 244 or
                              resid 300 or resid 304 or resid 305 or
                              resid 306 or resid 327 or resid 332 or
                              resid 334 or resid 367 or resid 368 or
                              resid 369 or resid 370) and
                              (name cb or name hb1 or name hb2 or
                              name sg1 or name sg2)'''))
    minc.setGroupList( groupList )
    protocol.initMinimize(minc)
    planeDistTools.topologySetup(minc,planePot)
    protocol.cartesianTopology(minc)

    # final torsion angle minimization
    #
    protocol.initMinimize(dyn, printInterval=50)
    dyn.run()

    # final all- atom minimization
    #
    protocol.initMinimize(minc, potList=potList, dEPred=10)
    minc.run()

    #do analysis and write structure
    loopInfo.writeStructure(potList)
    pass



from simulationTools import StructureLoop, FinalParams
StructureLoop(numStructures=numberOfStructures,
              pdbTemplate=outFilename,
              structLoopAction=calcOneStructure,
              genViolationStats=1,
              averagePotList=potList,
#              averageSortPots=[potList['BOND'],potList['ANGL'],potList['IMPR'],
#                               noe,rdcs,potList['CDIH']],
              averageSortPots=[potList['BOND'],potList['ANGL'],potList['IMPR']],
              averageTopFraction=0.5, #report only on best 50% of structs
              averageAccept=accept,   #only use structures which pass accept()
              averageContext=FinalParams(rampedParams),
              averageFilename="ave.pdb",    #generate regularized ave structure
              averageFitSel="name CA",
              averageCompSel="not resname ANI and not hydro").run()

