#!/usr/bin/python
import psfGen
#psfGen.pdbToPSF('bead1.pdb')

import protocol
#protocol.initRandomSeed(42)
#protocol.genExtendedStructure('extended.pdb')
opts,args=xplor.parseArguments(["inputfiles"])
inputPDB=args[0]
protocol.loadPDB(inputPDB)

import regularize
regularize.addUnknownAtoms_new()

#protocol.writePDB("allAtom.pdb")

#
from potList import PotList
potList = PotList()
crossTerms=PotList('cross terms')
numberOfStructures = 5
startStructure = 0
outFilename = "SCRIPT_%d_STRUCTURE_MEMBER.pdb" %  numberOfStructures
#
# Lists highTempParams and rampedParams will hold simulationTools.StaticRamp and
# simulationTools.MultRamp objects to handle parameter changes between the high
# termperature and annealing stage, and within the annealing stage (e.g., ramped
# force constants).
#
from simulationTools import StaticRamp, MultRamp, InitialParams, AnnealIVM

highTempParams = []
rampedParams = []


from globDiffPotTools import create_GlobDiffPot
#
globDiffA = create_GlobDiffPot('globDiffA',
                               selection="segid A",
                               globFilename=args[1],
                               useChainID=True,
                               globMap=[("name P","name P"),
                                        ("name C1' C2' C3' C4' O4'","name S"),
                                        ("resname ADE and name N1", "name Ab"),
                                        ("resname GUA and name N1", "name Gb"),
                                        ("resname CYT and name N3", "name Cb"),
                                        ("resname URI and name N3", "name Ub"),
                                       ])

globDiffA.setThreshold(0.5)
globDiffA.setUpperBound(0.1)
globDiffA.setScale(3)
potList.append( globDiffA )

globDiffB = create_GlobDiffPot('globDiffB',
                               selection="segid B",
                               globFilename=args[1],
                               useChainID=True,
                               globMap=[("name P","name P"),
                                        ("name C1' C2' C3' C4' O4'","name S"),
                                        ("resname ADE and name N1", "name Ab"),
                                        ("resname GUA and name N1", "name Gb"),
                                        ("resname CYT and name N3", "name Cb"),
                                        ("resname URI and name N3", "name Ub"),
                                        ])

globDiffB.setThreshold(0.5)
globDiffB.setUpperBound(0.1)
globDiffB.setScale(3)
potList.append( globDiffB )


                              
from regularize import fixupCovalentGeom
fixupCovalentGeom()
if xplor.p_processID==0:
    protocol.writePDB("start.pdb")
    pass
from regularize import fixupCovalentGeomIVM
#fixupCovalentGeomIVM(extraTerms=[globDiff],verbose=1,useVDW=True)
#fixupCovalentGeomIVM(verbose=1,useVDW=True) - maybe re-enable?

from posDiffPotTools import create_PosDiffPot
refRMSD = create_PosDiffPot("refRMSD","not name H*",
                            pdbFile=inputPDB,
                            cmpSel="not name H*")
refRMSD.setUpperBound(2)
refRMSD.setScale(10)
potList.append(refRMSD)


#print(xray.calcEnergy())
from solnScatPotTools import fitParams

#
# Set up distance restraint potential (e.g., from NOEs).
#
import noePotTools
noe = PotList('noe')
import genDist
basePairResids = genDist.readDuplexes(args[2])

restraintsA = genDist.genRestraints(basePairResids,
          selTemplate=r"segid A and resid %d and not name H* and not name *'")
potA = noePotTools.create_NOEPot('bpDistA', restraints=restraintsA)
noe.append(potA)

restraintsB = genDist.genRestraints(basePairResids,
          selTemplate=r"segid B and resid %d and not name H* and not name *'")
potB = noePotTools.create_NOEPot('bpDistB', restraints=restraintsB)
noe.append(potB)
potList.append(noe)
rampedParams.append(MultRamp(2, 30, "noe.setScale(VALUE)"))

#
# Set up potential for base-pair planarity restraints.
#
from xplorPot import XplorPot
#protocol.initPlanarity('plane.tbl')
#potList.append(XplorPot('PLAN'))
# (The setup of this term remains unchanged throughout; no need to involve
# highTempParams and/or rampedParams.)


#
# Set up potential for base-pair planarity restraints.
#
protocol.initPlanarity(basePairs=basePairResids)
potList.append(XplorPot('PLAN'))
# (The setup of this term remains unchanged throughout; no need to involve
# highTempParams and/or rampedParams.)


#
# Setup parameters for atom-atom repulsive term (van der Waals-like term).
#
from repelPotTools import create_RepelPot,initRepel
repel = create_RepelPot('repel')
potList.append(repel)
rampedParams.append( StaticRamp("initRepel(repel,use14=False)") )
rampedParams.append( MultRamp(.004, 4, "repel.setScale( VALUE)") )
# nonbonded interaction only between C1' atoms
highTempParams.append( StaticRamp("""initRepel(repel,
                                               use14=True,
                                               scale=0.004,
                                               repel=1.2,
                                               moveTol=45,
                                               interactingAtoms="name C1'"
                                               )""") )
#
# Set up statistical torsion angle potential (torsionDB).
#
from torsionDBPotTools import create_TorsionDBPot
torsiondb = create_TorsionDBPot(name='torsiondb',system='rna')
potList.append(torsiondb)
rampedParams.append(MultRamp(0.5, 4, "torsiondb.setScale(VALUE)"))


# Selected 1-4 interactions.
# ---
# 1-4 interactions are ommited from the above term (repel) because they are
# mostly affected by torsionDB (torsiondb).  However, torsionDB doesn't affect
# torsions of terminal, protonated groups (e.g., methyl rotation).  Thus, 1-4
# interactions involved in such torsions are set up here.  (Note: a more elegant
# solution might be implemented in the near future.)
#
from torsionDBPotTools import create_Terminal14Pot
repel14 = create_Terminal14Pot('repel14')
potList.append(repel14)
highTempParams.append(StaticRamp("repel14.setScale(0)"))
rampedParams.append(MultRamp(0.004, 4, "repel14.setScale(VALUE)"))



#
# Set up bond length potential.
# (Needed even if no Cartesian minimization is used, for "broken" rings.)
#
potList.append(XplorPot('BOND'))
# (The setup of this term remains unchanged throughout; no need to involve
# highTempParams and/or rampedParams.)


#
# Set up bond angle potential.
# (Needed even if no Cartesian minimization is used, for "broken" rings.)
#
potList.append(XplorPot('ANGL'))
rampedParams.append(MultRamp(0.4, 1.0, "potList['ANGL'].setScale(VALUE)"))

#
# Set up improper dihedral angle potential.
# (Needed even if no Cartesian minimization is used, for "broken" rings.)
#
potList.append(XplorPot('IMPR'))
rampedParams.append(MultRamp(0.1, 1.0, "potList['IMPR'].setScale(VALUE)"))


#
# Set up statistical base-base positional potential.

# In general, however, all RNA residues should be selected (e.g.,
# for an isolated RNA molecule use selection='all').
# Reference: Clore, GM & Kuszewski, J, (2003) J. Am. Chem. Soc. 125:1518-1525.
#
#FIX: add pairs?
protocol.initOrie(system='rna', selection='all',
                  #distCutoff=15, will speed things up considerably - supported
                  #               in Xplor-NIH versions 3.2.2 and up.
                  )

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

potList.append(XplorPot('HBON'))
potList['HBON'].setScale(1.0)

#
# Done with energy terms.

#
# Give atoms uniform weights, except for anisotropy axes (if any).
#
protocol.massSetup()


#
# Set up IVM object(s).
#

# IVM object for torsion-angle dynamics/minimization.
import ivm
dyn = ivm.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')

## Optional IVM object for final Cartesian minimization.
minc = ivm.IVM()
##
##for tensor in tensors.values():
##    tensor.setFreedom("varyDa, varyRh") # allow all tensor parameters float here
##
protocol.cartesianTopology(minc)



#
# Temperature set up.
#
temp_ini = 3000.0   # initial temperature
temp_fin = 25.0     # final temperature



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

    # 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=temp_ini,
                          initVelocities=1,
                          finalTime=150,    # stops at 150ps or 15000 steps
                          numSteps=15000,   # whichever comes first
                          printInterval=100)

    dyn.setETolerance( temp_ini/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=400,       #at each temp: 400 steps or
                          finalTime=1 ,       # 1 ps, whichever is less
                          printInterval=100)

    # perform simulated annealing
    #
    AnnealIVM(initTemp=temp_ini,
              finalTemp=temp_fin,
              tempStep=12.5,
              ivm=dyn,
              rampedParams=rampedParams).run()
              
    # 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, crossTerms)
    pass



from simulationTools import StructureLoop, FinalParams
StructureLoop(numStructures=numberOfStructures,
              pdbTemplate=outFilename,
              startStructure=startStructure,
              structLoopAction=calcOneStructure,
              genViolationStats=1,
              averagePotList=potList,
#              averageSortPots=[potList['BOND'],potList['ANGL'],potList['IMPR'],
#                               noe,rdcs,potList['CDIH']],
              averageCrossTerms=crossTerms,
              averageTopFraction=0.2, #report only on best 50% of structs
              averageContext=FinalParams(rampedParams),
              averageFilename="SCRIPT_ave.pdb",    #generate regularized ave structure
              averageFitSel="name P",
              averageCompSel="not PSEUDO and not name H*"     ).run()

