# ----- sample of docking script with diffusion temperature optimization ----- #
# ---- which uses NMR relaxation data and chemical shift pertrubation data --- #
# ------------------------- as structural restraints ------------------------- # 
# ----------- for reference see Ryabov, Clore, and Schwieters (2010) --------- #

# this checks for typos on the command-line. User-customized arguments can
# also be specified.
#
(opts,args) = xplor.parseArguments(["quick"])

quick=False
for opt in opts:
    if opt[0]=="quick":  #specify -quick to just test that the script runs
        quick=True
        pass
    pass

outFilename = "SCRIPT_STRUCTURE.sa"
numberOfStructures = 2 if quick else 512 
numberOfLoops      = 2 if quick else 50
numberOf_cool_Loops      = 1 


# protocol module has many high-level helper functions. #
import protocol

# explicitly set random seed #
protocol.initRandomSeed(7654)

command = xplor.command

protocol.initParams("protein")

# read in the EIN domain coordinates #
protocol.loadPDB("randomized_sch_EIN.pdb") 

# read in the HPr domain coordinates #
protocol.loadPDB("randomized_sch_HPr.pdb") 

xplor.simulation.deleteAtoms("not known")

# subtract EIN center of gravity coordinates #
# from original EIN coordinates #
xplor.command("vector do (x=x+6.6) (resid 1:250)")
xplor.command("vector do (y=y-2.3) (resid 1:250)")
xplor.command("vector do (z=z-138.5) (resid 1:250)") 

# subtract HPr center of gravity coordinates #
# from original HPr coordinates #
xplor.command("vector do (x=x-13) (resid 301:385)")
xplor.command("vector do (y=y-7.9) (resid 301:385)")
xplor.command("vector do (z=z-8.8) (resid 301:385)") 

# create PotLists which contains lists of potential terms. #

from potList import PotList

potList = PotList()  
score = PotList()  
potList_0 = PotList()  
potList_1 = PotList() 
potList_2 = PotList() 


from simulationTools import MultRamp, StaticRamp, InitialParams

rampedParams=[]
highTempParams=[]


from xplorPot import XplorPot


# conract shifts
noe=PotList('noe')
potList_0.append(noe)
score.append(noe)
potList.append(noe)
potList_1.append(noe)
potList_2.append(noe)

from noePotTools import create_NOEPot
for (name,scale,file) in [('all',10,"shifts_noe_newx_EINHPr.tbl")]:
    pot = create_NOEPot(name,file)
    pot.setScale(scale)
    noe.append(pot)
rampedParams.append( MultRamp(0.3,60, "noe.setScale( VALUE )") )

# set up Miyazawa contact potential
from residueAffPotTools import create_ResidueAffPot 
ra = create_ResidueAffPot('hydphob', interdomainContacts=True, 
                          potentialName="Miyazawa")
ra.setAveType("center")
ra.setMoveTol(0.5)
ra.setCutoffLong(20) 
ra.setScale(10)

potList_1.append(ra)
potList_2.append(ra)

rampedParams.append( MultRamp(1,50,"ra.setScale(VALUE)") )

# IVM setup
#   the IVM is used for performing dynamics and minimization in torsion-angle
#   space, and in Cartesian space.
#
from ivm import IVM
from ivm import PublicIVM
dyn_fix = IVM()		# ivm for rigid body dynamics
dyn_free_sch = IVM()

# reset ivm topology for torsion-angle dynamics
dyn_fix.reset()
dyn_free_sch.reset()

protocol.torsionTopology(dyn_fix)
protocol.torsionTopology(dyn_free_sch)

# ivms used for initial rigid body gradient minimization #
dyn_fix.group( 'resid 1:249' )
dyn_fix.group( 'resid 301:385' )

# ivms used for simulated annealing which let side chains to be flexible  #
dyn_free_sch.group( 'resid 1:249 and \
                     (name CA or name C or name N or name O or name HN)' )
dyn_free_sch.group( 'resid 301:385 and \
                     (name CA or name C or name N or name O or name HN)' )

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


# compare atomic Cartesian rmsd with a reference structure
#  backbone and heavy atom RMSDs will be printed in the output
#  structure files

from posDiffPotTools import create_PosDiffPot

refRMSD = create_PosDiffPot("refRMSD",
          selection="(name CA and resid 3:249) or (name CA and resid 301:385)", 
          selection2="(name CA and resid 3:249) or (name CA and resid 301:385)",
          pdbFile='reference_docking_structure.pdb',
          cmpSel=" (name CA and resid 3:249) or (name CA and resid 301:385) ")

# setup parameters for atom-atom repulsive term. (van der Waals-like term)
#
potList_1.append( XplorPot('VDW') )
potList_2.append( XplorPot('VDW') )
score.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(.01,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,
                                                        rcon=0.01,
                                                        onlyCA=1)""") )

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


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




# object which performs simulated annealing
#
from simulationTools import AnnealIVM
init_t  = 1000 # 3000.     # Need high temp and slow annealing to converge
cool = AnnealIVM(initTemp =init_t,
                 finalTemp=10,
                 tempStep =100 if quick else 10,
                 ivm=dyn_free_sch,
                 rampedParams = rampedParams)


from atomAction import randomizeDomainPos

import random

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

    initial_tmp_pos = xplor.simulation.atomPosArr()
    tmp_pos_swap = xplor.simulation.atomPosArr()

    tmp_energy = 1e9	# big number
    k=0			# set up initial minimaizing loop

    while k < numberOfLoops:

    	xplor.simulation.setAtomPosArr(initial_tmp_pos)

    	randomizeDomainPos( 'resid 1:249', deltaPos=45 )


    	# initialize parameters for high temp dynamics. #
    	InitialParams( rampedParams )
    	InitialParams( highTempParams )

    	protocol.initMinimize(dyn_fix,potList=potList_0,
                          printInterval=50)
    	dyn_fix.run()

    	dyn_fix.run()

    	protocol.initMinimize(dyn_fix,potList=potList,
                          printInterval=50)
    	dyn_fix.run()

    	dyn_fix.run()

    	protocol.initMinimize(dyn_fix,potList=score,
                          printInterval=50)
    	dyn_fix.run()

    	dyn_fix.run()

    	dyn_fix.run()

	print potList.calcEnergy(), score['VDW'].calcEnergy() 

	tmp_energy_swap = score.calcEnergy()

	print tmp_energy_swap

	if tmp_energy_swap < tmp_energy :
		tmp_pos_swap = xplor.simulation.atomPosArr()
		tmp_energy = tmp_energy_swap

	k=k+1
	print k

    xplor.simulation.setAtomPosArr(tmp_pos_swap)

    tmp_energy = potList.calcEnergy() # debugging line
    print tmp_energy		      # debugging line



    k2=0			# set up initial cooling loop
    while k2 < numberOf_cool_Loops:

    	InitialParams( rampedParams )

    	# initialize integrator for simulated annealing
    	protocol.initDynamics(dyn_free_sch,
                              potList=potList_2,
                              numSteps=3 if quick else 600,
                              finalTime=1 ,
                              printInterval=100)

    	# perform simulated annealing
    	cool.run()
         
	k2=k2+1
	print k2
    
              
    # final torsion angle minimization
    protocol.initMinimize(dyn_free_sch,potList=potList_2,
                          printInterval=50)
    dyn_free_sch.run()

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


from simulationTools import StructureLoop, FinalParams
StructureLoop(numStructures = numberOfStructures , 
              pdbTemplate=outFilename,
              structLoopAction=calcOneStructure,
              genViolationStats=1,
              averagePotList=potList_2,
              averageSortPots=[noe, potList_2['ANGL'],
                               potList_2['BOND'], potList_2['IMPR'], 
                               potList_2['RAMA'], potList_2['VDW'], ra],
              averageCrossTerms=refRMSD,
              averageTopFraction=0.0195, 
              averageContext=FinalParams(rampedParams),
              averageFilename="SCRIPT_ave.pdb",    #generate regularized ave structure
              averageFitSel="name CA",
              averageCompSel="not resname ANI and not hydro" ).run()

