#calculates starting structure in Xplor-NIH with a simulated annealing approach
#automatically detects if ligand is involved, but ligands are handled as rigid bodies
#based on different Xplor tutorial script

#parent script: preparing_starting_structures.sh
#09/14/19, JD




xplor.requireVersion("3.8")

# 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 = '../start_structures/SCRIPT_STRUCTURE.pdb'


numberOfStructures=200   #is automatically changed to 200 if selection of 100 best structures is performed



import protocol


#decide if drug is included or not for parameter loading. Also determine the residue number of the 3'-end and numbers for the ligands
with open('../sequence', 'r') as seqFile:
    sequence = [row for row in seqFile]
seq_length = len(sequence)

try:
    with open('../LIG.psf', 'r') as fh:
        protocol.initParams(files=['../LIG.par','nucleic'])
        command = xplor.command 
        protocol.initStruct(['../sequence.psf','../LIG.psf'])      
        lig_1 = 'resid ' + str(seq_length)
        seq_extend = 'resid 1:' + str(seq_length-1)

except IOError:
    protocol.initParams(files=['nucleic'])
    command = xplor.command 
    protocol.initStruct(['../sequence.psf'])     
    seq_extend = 'resid 1:' + str(seq_length)

# Generate extended conformation with satisfied covalent geometry.
#
protocol.initCoords("../sequence.pdb",useChainID=True,includeHETATM=True) 


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

from xplorPot import XplorPot

# Parameters to ramp up during the simulated annealing protocol
from simulationTools import MultRamp, StaticRamp, InitialParams
rampedParams=[]
highTempParams=[]

# Set up NOE and hydrogen bond potential
noe=PotList('noe')
potList.append(noe)
from noePotTools import create_NOEPot

for (name,scale,filename) in [('noe', 1, '../restraints/distance.xplor'),
                               ('hbond', 100, '../restraints/hbond.tbl'),
                          #add entries for additional tables
                          ]:
    pot = create_NOEPot(name,filename)
    pot.setPotType("soft") #- if you think there may be bad NOEs
    pot.setScale(scale)
    pot.setThreshold(0.1)
    noe.append(pot)

rampedParams.append( MultRamp(20,60, "noe.setScale( VALUE )") )

# Set up dihedral angles
from xplorPot import XplorPot
protocol.initDihedrals('../restraints/dihedral.tbl',
                       #useDefaults=False  # by default, symmetric sidechain
                                           # restraints are included
                       )
potList.append( XplorPot('CDIH') )
highTempParams.append( StaticRamp("potList['CDIH'].setScale(200)") )
rampedParams.append( StaticRamp("potList['CDIH'].setScale(200)") )
potList['CDIH'].setThreshold( 5 ) # Set threshold to 5 degrees 

# Use statistical torsion angle potential
# select residues in G-core (or duplex)
from torsionDBPotTools import create_TorsionDBPot
torsiondb = create_TorsionDBPot('torsiondb', system="dna",selection='resid 1:6 or resid 25:30')
potList.append(torsiondb)
rampedParams.append( MultRamp(.004,1,"torsiondb.setScale(VALUE)") ) 


#
# Set up potential for base-pair planarity restraints.
#
protocol.initPlanarity('../restraints/plane.tbl')
potList.append(XplorPot('PLAN'))
highTempParams.append( StaticRamp("potList['PLAN'].setScale(0.1)") )
rampedParams.append( MultRamp(0.1, 1, "potList['PLAN'].setScale( VALUE)"))


#
# Setup parameters for atom-atom repulsive term (van der Waals-like term).
#
#=>>>>>>> select residues in G-core (or duplex)
from repelPotTools import create_RepelPot,initRepel
repel = create_RepelPot('VDW')
potList.append(repel)
rampedParams.append( StaticRamp("initRepel(repel,use14=False)") )
rampedParams.append( MultRamp(0.001, 4, "repel.setScale( VALUE)") )

highTempParams.append( StaticRamp("""initRepel(repel,
                                               use14=True,
                                               scale=0.00001, 
                                               repel=1.2,
                                               moveTol=45,
                                               interactingAtoms="(name C2 or name C4 or name C5 or name C6 or name N1 or name N3 or name N7 or name N9 or name P) and (resid 1:6 or resid 25:30)"
                                               )""") )

try:
    with open('../LIG.psf', 'r') as fh:
        highTempParams.append( StaticRamp("""initRepel(repel,
                                               use14=False,
                                               scale=0.05, 
                                               repel=0.9,
                                               moveTol=45,
                                               interactingAtoms="resname LIG"
                                               )""") )    
except IOError:
    pass

#
# Set up statistical base-base positional potential.
# ---
# select residues in G-core (or duplex)
protocol.initOrie(system='dna', selection='resid 1:6 or resid 25:30') 

potList.append(XplorPot('ORIE'))
rampedParams.append(MultRamp(0.002,1,"xplor.command('orie scale VALUE end')"))



# Enforce geometry (bonds, angles, and impropers)
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
protocol.massSetup()



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


# Initially 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") )


try:
    with open('../LIG.psf', 'r') as fh:
        dyn.breakAllBondsIn("not resname LIG")     
except IOError:
    pass
    # do nothing w/o ligands


#protocol.initMinimize(dyn,numSteps=200)

dyn.run() 

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

protocol.torsionTopology(dyn) 


# minc used for final cartesian minimization
#
minc = IVM()

# Argument flexRiboseRing below is a string that selects residues whose ribose
# rings will have all endocyclic angles flexible.  In general, all ribose rings
# should be selected.  In this example, the non-RNA ligand (residue 34) has to
# be excluded.
#protocol.torsionTopology(dyn, flexRiboseRing='all')


protocol.initMinimize(minc)

try:
    with open('../LIG.psf', 'r') as fh:
        minc.group(lig_1)     
except IOError:
    pass
    # do nothing w/o ligand

protocol.cartesianTopology(minc)

# Create object that performs simulated annealing
#
from simulationTools import AnnealIVM
init_t  = 10000
cool = AnnealIVM(initTemp =init_t,
                 finalTemp=20,
                 tempStep =20,
                 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.
    """

    # Generate initial structure by randomizing torsion angles.
    import monteCarlo
    monteCarlo.randomizeTorsions(dyn)

    # Set torsion angles from restraints.
    # (They start satisfied, allowing the shortening of high temp dynamics.)
    import torsionTools
    torsionTools.setTorsionsFromTable('../restraints/dihedral.tbl')
    try:
        with open('../LIG.psf', 'r') as fh:
            dyn.group(lig_1)
    except IOError:
        pass
        # do nothing w/o ligand 


    # Fix up covalent geometry.
    # (The torsion restraints may include ring torsions and distort geometry.)
    while True:
       try:
            protocol.fixupCovalentGeom(maxIters=300, useVDW=1,sel= seq_extend)
            break
       except protocol.CovalentViolation:
            pass


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



    # high temp dynamics
    #
    protocol.initDynamics(dyn,
                          potList=potList, # potential terms to use
                          bathTemp=init_t,
                          initVelocities=0, # 0 should be randomized
                          finalTime=30,    # stops at 15 ps or 5000 steps
                          numSteps=10000,   # whichever comes first
                          printInterval=1000)

    dyn.setETolerance( init_t/100 )  #used to det. stepsize. default: t/1000
    try:
        with open('../LIG.psf', 'r') as fh:
            dyn.group(lig_1)
    except:
        pass
    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=0.2 ,       # 0.2 ps, whichever is less
                          printInterval=100)

    # Perform simulated annealing
    #
    cool.run()

    # Final torsion angle minimization
    #
    protocol.initMinimize(dyn,
                          printInterval=100)
    dyn.run()

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

    #do analysis and write structure when this routine is finished.
    pass


from simulationTools import StructureLoop, FinalParams
StructureLoop(numStructures=numberOfStructures,
              pdbTemplate=outFilename,
              structLoopAction=calcOneStructure,
              doWriteStructures=True,     # analyze and write coords after calc
              genViolationStats=True,
             # averageTopFraction=0.2,    #report stats on best 50% of structs
              #averageContext=FinalParams(rampedParams),
              averagePotList=potList
              ).run()
