
xplor.requireVersion("2.39")


# 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


# filename for output structures. This string must contain the STRUCTURE
# literal so that each calculated structure has a unique name. The SCRIPT
# literal is replaced by this filename (or stdin if redirected using <),
# but it is optional.
#
outFilename = "SCRIPT_STRUCTURE.sa"
numberOfStructures=2 if quick else 64

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

protocol.initRandomSeed()   #set random seed - by time

# parameters
protocol.initParams(("protein","nucleic"))

# psf
protocol.initStruct("inputs/complex.psf")

# coords
protocol.initCoords("inputs/complex.pdb")
xplor.simulation.deleteAtoms("resname ANI")

import protocol
protocol.fixupCovalentGeom(useVDW=True,verbose=1,
                           angleTol=2.5)


nconf = 1
import protocol, xplor
#
# a PotList contains a list of potential terms. This is used to specify which
# terms are active during refinement.
#
from potList import PotList
potList = PotList()
crossTerms = PotList()

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

rampedParams=[]
highTempParams=[]

# 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","name CA or name C or name N",
#                            pdbFile='g_xray.pdb',
#                            cmpSel="not name H*")
#crossTerms.append(refRMSD)

# orientation Tensor - used with the dipolar coupling term
#  one for each medium
#   For each medium, specify a name, and initial values of Da, Rh.
#

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

# add clock pseudo-atoms
#
from prePotTools import addClockAtoms
for resid in range(777,783):
   addClockAtoms(resid)
   pass


#
# for Es2
#

# for Es4
#


es4at800 = PotList('es4at800')
for (name,file,assignType) in [
    ("Es4hn800","preG2hn800.tbl","normal"),]:
    p = create_PREPot(name,"inputs/"+file,assignType,800,clockResid=777)
    es4at800.append(p)
    pass
pre.append(es4at800)

def runSBmode():
    print 'configuring SB mode'
    from prePotTools import setupSBmode
    from simulationTools import flattenPotList
    for p in flattenPotList(pre): setupSBmode(p)
    #apparent tauc
#    tcappEs2=2.9
    tcappEs4=13.0
#    for p in list(es2at500)+list(es3at500):
#        p.setTcType("fix")
#        p.setTauC( tcappEs2 )
#        pass
    for p in es4at800:
        p.setTcType("fix")
        p.setTauC( tcappEs4 )
        pass
#    for p in list(es2at750)+list(es3at750):
#        p.setTcType("opt")
#        p.setTauC(  tcappEs2       )
#        p.setTcMin( tcappEs2       )
#        p.setTcMax( tcappEs2 * 1.6 )
#    for p in es4at800:
#        p.setTcType("opt")
#        p.setTauC(  tcappEs4       )
#        p.setTcMin( tcappEs4       )
#        p.setTcMax( tcappEs4 * 1.6 )
#        pass
    return

bbPre = PotList('bbPre')
for p in [es4at800['Es4hn800']]:
    bbPre.append(p)
    pass
crossTerms.append(bbPre)



def runSBMFmode():
    print 'configuring SBMF mode'
    from prePotTools import setupSBMFmode
    from simulationTools import flattenPotList
    for p in flattenPotList(pre): setupSBMFmode(p)
    pre.calcEnergy()
#    tcappEs2=2.9
    tcappEs4=13.0
#    maxtcEs2 = tcappEs2 / es2at500['Es2hn500'].aveS2()
    maxtcEs4 = tcappEs4 / es4at800['Es4hn800'].aveS2()

    for p in es4at800 :
        p.setTcType("opt")
        p.setTcMin( tcappEs4 )
        p.setTcMax( maxtcEs4 )
        pass

    return

potList.append(pre)
rampedParams.append( MultRamp(0.05,1.0, "pre.setScale( VALUE )") )
rampedParams.append( StaticRamp("runSBMFmode()") )
highTempParams.append( StaticRamp("runSBmode()") )

# set up NOE potential
noe=PotList('noe')
potList.append(noe)
from noePotTools import create_NOEPot
for (name,scale,file) in [('noe',1,["inputs/noe_prot.tbl",
                                    "inputs/noe_300.tbl",
                                   #"inputs/6_30n15noe_arg.tbl",
                                    #"inputs/7_2metonly.tbl",
                                    #"inputs/helix_hbond.tbl",
                                    #"inputs/dna_trial1-new.tbl", # modified version of dna_trial1.tbl 
                                    #"inputs/endnoe_dnax.tbl",
                                    #"inputs/pp_trial.tbl",
                                    #"inputs/trial_inter3.tbl",
                                    #"inputs/6_22nh_inter.tbl",
                                    #"inputs/inter_hb.tbl",
                                    #"inputs/ends_bdnazzz-new.tbl", # modified version of ends_bdnazzz.tbl
                                    ]),
                          ('hbond',1,["inputs/hbond_pro.tbl",
                                      "inputs/hbond_300.tbl", 
                                      ])
                          ]:
    pot = create_NOEPot("noe_"+name,file)
    # pot.setPotType("soft") # if you think there may be bad NOEs
    pot.setScale(scale)
    noe.append(pot)
rampedParams.append( MultRamp(2,30, "noe.setScale( VALUE )") )

from xplorPot import XplorPot
# carbon 13 chemical shifts
#protocol.initCarb13("inputs/c13_shifts.tbl")
#potList.append(XplorPot("CARB"))

# set up J coupling - with Karplus coefficients
#from jCoupPotTools import create_JCoupPot
#jCoup = create_JCoupPot("jcoup","inputs/hnha_xplr.tbl",
#                        A=6.98,B=-1.38,C=1.72,phase=-60.0)
#potList.append(jCoup)
#x_jCoup = create_JCoupPot("xjcoup","inputs/dna_pcoup.tbl",
#                          A=15.3,B=-6.2,C=1.5,phase=120.0)
#crossTerms.append(x_jCoup)

# Set up dihedral angles
protocol.initDihedrals(["inputs/dihedrals_prot.tbl",
                        "inputs/dihe_dna.tbl",])
                        #"inputs/dihed_dna_edtaBform.tbl",
                        #"inputs/dihed_dna_edtaBformEs3.tbl",
                        #"inputs/chi_final-altered.tbl"],
                        #useDefaults=0)
potList.append( XplorPot('CDIH') )
highTempParams.append( StaticRamp("potList['CDIH'].setScale(10)") )
rampedParams.append( StaticRamp("potList['CDIH'].setScale(200)") )
# set custom values of threshold values for violation calculation
#
potList['CDIH'].setThreshold( 5 )



# radius of gyration term 
#
protocol.initCollapse("resid 4:81 or resid 105:112 or resid 117:124",
                      Rtarget=15,scale=0.25)
potList.append( XplorPot('COLL') )

#Rama torsion angle database
#
protocol.initRamaDatabase('nucleic')
protocol.initRamaDatabase('protein')
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'] 

#FIXME: be able to tune nucleic fc separate from protein
rampedParams.append( MultRamp(1,1,"potList['RAMA'].setScale(VALUE)") )
rampedParams.append( MultRamp(.002,1,
 """for pclass in ramaProtClasses:
      xplor.command('rama class %s force %f end ' % (pclass, VALUE))""") )

#initialize the aa-aa positional database
xplor.command("@inputs/bbpd.setup")
potList.append( XplorPot('ORIE') )

#planarity restraints
xplor.command("restraints plane @inputs/plane_dna.tbl end")
potList.append(XplorPot("plan"))

#
# setup parameters for atom-atom repulsive term. (van der Waals-like term)
#
potList.append( XplorPot('VDW') )
rampedParams.append( StaticRamp("protocol.initNBond()") )
rampedParams.append( StaticRamp("setConstraints()") )

rampedParams.append( MultRamp(0.9,0.78,
                         "xplor.command('param nbonds repel VALUE end end')") )
rampedParams.append( MultRamp(.004,4,
                         "xplor.command('param nbonds rcon VALUE end end')") )
# nonbonded interaction only between CA atoms
highTempParams.append( StaticRamp("""protocol.initNBond(cutnb=100,
                                                        rcon=1,
                                                        tolerance=45,
                                                        repel=1.2,
                                                        onlyCA=1)""") )
highTempParams.append(
    StaticRamp(
    xplor.command("""constraints 
       inter = (not (name ca or name p or name mn)) (all) 
       weights * 1 angl 0.4 impr 0.1 vdw 0 elec 0 end 
       inter = (name ca or name p or name mn) (name ca or name p or name mn) 
       weights * 1 angl 0.4 impr 0.1 vdw 1.0 end 
      end""") ) )

def setConstraints():
    #number of conformers (support up to 10 conformers)
    scaleVdw = 1.0 / float( nconf )
    xplor.command("""
      constraints 
      !   interaction (all) (all) 
         inter = (segid " ") (segid " ")
         inter = (segid "ALT1") (segid "ALT1")
         inter = (segid "ALT2") (segid "ALT2")
         inter = (segid "ALT3") (segid "ALT3")
         inter = (segid "ALT4") (segid "ALT4")
         inter = (segid "ALT5") (segid "ALT5")
         inter = (segid "ALT6") (segid "ALT6")
         inter = (segid "ALT7") (segid "ALT7")
         inter = (segid "ALT8") (segid "ALT8")
         inter = (segid "ALT9") (segid "ALT9")
         inter = (segid "ALT0") (segid "ALT0")
         weights * 1 end 
         inter = (segid " ") (segid "ALT1")
         inter = (segid " ") (segid "ALT2")
         inter = (segid " ") (segid "ALT3")
         inter = (segid " ") (segid "ALT4")
         inter = (segid " ") (segid "ALT5")
         inter = (segid " ") (segid "ALT6")
         inter = (segid " ") (segid "ALT7")
         inter = (segid " ") (segid "ALT8")
         inter = (segid " ") (segid "ALT9")
         inter = (segid " ") (segid "ALT0")
         weights * 1 vdw %f end 
      end""" %  scaleVdw )
    return



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
#AtomSel("not resname ANI").apply( SetProperty("mass",100.) )
#varTensorTools.massSetup(media.values(),300)
AtomSel("resname TAU ").apply(SetProperty("mass", 300.0))
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()

# group EDTA-Mn
for res in (303, 316) :
   for i in range(0,nconf+1) :
      dyn.group("""((resid %d and (name OY* or name NY* or name HY* or name CY*))
               or resid %d)""" % (i, res[0], res[1]))
      pass
   pass

# allow fo tau_c optimization
from simulationTools import flattenPotList
for p in flattenPotList(pre): 
   p.setTcType("opt")
   pass



# initialize ivm topology for torsion-angle dynamics

#for m in media.values():
#    m.setFreedom("fixDa, fixRh")        #fix tensor Rh, Da, vary orientation
#    m.setFreedom("varyDa, varyRh")      #vary tensor Rh, Da, vary orientation
protocol.torsionTopology(dyn)

# minc used for final cartesian minimization
#
minc = IVM()
protocol.initMinimize(minc)

# group EDTA-Mn
for res in ([ (303, 303), (316,316)]) :
   for i in range(0,nconf+1) :
      minc.group("""(resid %d and (name OY* or name NY* or name HY* or name CY*))
               or resid %d)""" % (i, res[0], res[1]))
      pass
   pass

#for m in media.values():
#    m.setFreedom("varyDa, varyRh")    #allow all tensor parameters float here
#    pass
protocol.cartesianTopology(minc)



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

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

    # calc. initial tensor orientation
    #
    from varTensorTools import calcTensorOrientation
   # for medium in media.values():
   #     calcTensorOrientation(medium)
   #    pass

    # initialize parameters for high temp dynamics.
    InitialParams( rampedParams )
    # high-temp dynamics setup - only need to specify parameters which
    #   differfrom initial values in rampedParams
    InitialParams( highTempParams )

    # high temp dynamics
    #
    protocol.initDynamics(dyn,
                          potList=potList, # potential terms to use
                          bathTemp=init_t,
                          initVelocities=1,
                          finalTime=10,   # stops at 800ps or 8000 steps
                          numSteps=5 if quick else 500, # whichever comes first
                          printInterval=100)

    dyn.setETolerance( init_t/100 )  #used to det. stepsize. default: t/1000

    dyn.run()

    # initialize parameters for cooling loop
    InitialParams( rampedParams )


    # initialize integrator for simulated annealing
    #
    protocol.initDynamics(dyn,
                          potList=potList,
                          numSteps=100,       #at each temp: 100 steps or
                          finalTime=.2 ,       # .2ps, whichever is less
                          printInterval=100)

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

    # optional cooling in Cartesian coordinates
    #
    protocol.initDynamics(minc,
                          potList=potList,
                          numSteps=100,       #at each temp: 100 steps or
                          finalTime=.4 ,       # .2ps, whichever is less
                          printInterval=100)

    protocol.initMinimize(minc,
                          potList=potList,
                          dEPred=10)
    minc.run()

    # print close-contact atoms to a .bumps file
    # this XPLOR command will honor the contraints interaction setup
    # so that overlaps the copies of the tag(s) are not counted. The
    # output in the .viols will report these overlaps.
    xplor.command("""
    distance
      from=(not PSEUDO) to=(known)
      cuton=0 cutoff=1.5 disp=print
      output=%s
    end
    """ % (loopInfo.filename() + '.bumps') )
    #do analysis and write structure
    loopInfo.writeStructure(potList)
    pass



from simulationTools import StructureLoop, FinalParams
StructureLoop(numStructures=numberOfStructures,
              pdbTemplate=outFilename,
              structLoopAction=calcOneStructure,
              genViolationStats=1,
              averageTopFraction=0.5, #report stats on best 50% of structs
              averageContext=FinalParams(rampedParams),
              averageCrossTerms=crossTerms,
              averageSortPots=[potList['BOND'],potList['ANGL'],potList['IMPR'],
                               noe,potList['CDIH']],
              averagePotList=potList).run()
