# ----- sample of docking script with diffusion temperature optimization ----- #
# --- for reference see Ryabov, Suh, Grishaev, Clore, and Schwieters (2009) -- #

# 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","nucleic"))

# or read initial PDB

protocol.loadPDB("inputs/complex.pdb") 

xplor.simulation.deleteAtoms("not known")

# move HPr domain close to the EIN center of gravity

xplor.command("vector do (x=x-15) (resid 94:207)")
xplor.command("vector do (y=y+20) (resid 94:207)")
xplor.command("vector do (z=z-10) (resid 94:207)") 


# a PotLists contains lists of potential terms. #

from potList import PotList

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

# parameters to ramp up during the simulated annealing protocol #

from simulationTools import MultRamp, StaticRamp, InitialParams

rampedParams=[]
highTempParams=[]


from xplorPot import XplorPot

# PRE restraints
#
from prePotTools import create_PREPot
pre = PotList('pre')

from prePotTools import create_PREPot
pre = PotList('pre')
for (name,file,assignType) in [
    ("Es4hn800","preG2hn800.tbl","normal"),]:
    p = create_PREPot(name,"inputs/"+file,assignType,800,tauc=13.0)
    pre.append(p)
    pass

potList_0.append(pre)
score.append(pre)
potList.append(pre)				
potList_1.append(pre)
potList_2.append(pre)
rampedParams.append( MultRamp(0.05,1.0, "pre.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) # potential is 0 for larger distance ra.setCutoffShort(6.5)
                     # potential has full value for smaller distance
ra.setScale(10)      # ``force constant''

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)

dyn_fix.group( 'resid 301:334' )
dyn_fix.fix( 'resid 94:207' )

dyn_free_sch.group( 'resid 301:334' )
dyn_free_sch.fix( 'resid 301:385 and (name CA or name C or name N or name O or name HN)' )

dyn_free_sch.fix( 'resid 1700' )
dyn_fix.fix( 'resid 1700' )

# 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.tbl")]:
#    pot = create_NOEPot(name,file)
#    pot.setScale(scale)
#    noe.append(pot)
#rampedParams.append( MultRamp(0.3,60, "noe.setScale( VALUE )") )




# 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,
                                                        onlyCA=1)""") )

#Rama torsion angle database
#
protocol.initRamaDatabase()
protocol.initRamaDatabase('nucleic')
potList_2.append( XplorPot('RAMA') )
potList.append( XplorPot('RAMA') )
ramaProtClasses=['phi_psi_ala', 'phi_psi_arg', 'phi_psi_asn', 'phi_psi_asp',
                 'phi_psi_chi1_arg', 'phi_psi_chi1_asn', 'phi_psi_chi1_asp',
                 'phi_psi_chi1_cyo', 'phi_psi_chi1_cyr', 'phi_psi_chi1_gln',
                 'phi_psi_chi1_glu', 'phi_psi_chi1_his', 'phi_psi_chi1_ile',
                 'phi_psi_chi1_leu', 'phi_psi_chi1_lys', 'phi_psi_chi1_met',
                 'phi_psi_chi1_phe', 'phi_psi_chi1_ser', 'phi_psi_chi1_thr',
                 'phi_psi_chi1_trp', 'phi_psi_chi1_tyr', 'phi_psi_chi1_val',
                 'phi_psi_csp', 'phi_psi_cyo', 'phi_psi_cyr', 'phi_psi_gln',
                 'phi_psi_glu', 'phi_psi_gly', 'phi_psi_his', 'phi_psi_ile',
                 'phi_psi_leu', 'phi_psi_lys', 'phi_psi_met', 'phi_psi_phe',
                 'phi_psi_pro', 'phi_psi_ser', 'phi_psi_thr', 'phi_psi_trp',
                 'phi_psi_tyr', 'phi_psi_val', 'phi_psi_xcp', 'phi_psi_xpr',
                 'chi1_chi2_arg', 'chi1_chi2_asn', 'chi1_chi2_asp',
                 'chi1_chi2_chi3_arg', 'chi1_chi2_chi3_gln',
                 'chi1_chi2_chi3_glu', 'chi1_chi2_chi3_lys',
                 'chi1_chi2_chi3_met', 'chi1_chi2_cyo', 'chi1_chi2_gln',
                 'chi1_chi2_glu', 'chi1_chi2_his', 'chi1_chi2_ile',
                 'chi1_chi2_leu', 'chi1_chi2_lys', 'chi1_chi2_met',
                 'chi1_chi2_phe', 'chi1_chi2_trp', 'chi1_chi2_tyr',
                 'chi2_chi3_arg', 'chi2_chi3_gln', 'chi2_chi3_glu',
                 'chi2_chi3_lys', 'chi2_chi3_met', 'chi4_arg', 'chi4_lys'] 
		 
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)") )


# radius of gyration term.
#
protocol.initCollapse("resid 301:334",
                      Rtarget=20.0)
potList_2.append( XplorPot('COLL') )
rampedParams.append( MultRamp(0.0005,5,"potList_2['COLL'].setScale(VALUE)") )



# object which performs simulated annealing
#
from simulationTools import AnnealIVM
init_t  = 500 # 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 301:334', deltaPos=45 )

    	# randomize temperature for the diffusion tensor settings
    	#rand_temp=random.uniform(dif_f.getMedianTmp()-dif_f.rangeTmpFit(), 
        #                        dif_f.getMedianTmp()+dif_f.rangeTmpFit())        

    	#reset_DiffPot_temp(dif_f,rand_temp)

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

    	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()


	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 300,# at each temp: 300 steps or
                              finalTime=0.5 ,        # .5ps, whichever is less
                              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=[potList_2['ANGL'], potList_2['BOND'], 
                               potList_2['IMPR'], potList_2['RAMA'], potList_2['VDW'], 
                               ra, potList_2['COLL']],
#              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()

