#!/usr/bin/env python

__author__ = "Randall J. Radmer"
__version__ = "1.0"
__doc__ = """Script to run simple amber simulation."""


#Simulation options
stepSize=0.001      #ps
stepsPerReport=100
numSteps=10*stepsPerReport
#temperature=300.0
temperature=0.0
#randomSeed=None
randomSeed=1
cutoff=12.0
verbose=True
useGBSA=False     # Don't use GBSA yet.  Still coding...
tryUsingCuda=True

import  os, sys, math

import simtk.utils.amber.amberOpenmmSystem as amberOpenmmSystem
import simtk.utils.amber.amberOpenmmContext as amberOpenmmContext
import simtk.utils.pdbWriter as pdbWriter


def main(verbose=False):
    """Main loop, builds an OpenMM context, adds particles and parameters,
       runs and dumps coords and forces"""
    import time
    time0=time.time()
    timeX=time.time()

    (args_proper,
     pdbInFilename,
     paramFilename,
     prepFilename,
     ctPrepFilename,
     ntPrepFilename,
     arMappingFilename,
     pdbOutFilename) = parseCommandLine()
    if (not pdbInFilename or
        not paramFilename or
        not prepFilename or
        not pdbOutFilename):
         usageError()

    if verbose:
        sys.stdout.write("Read pdb and data files, and setup simulation.\n")
        sys.stdout.flush()
    amberSystem=amberOpenmmSystem.AmberSystem(pdbInFilename,
                                              paramFilename,
                                              prepFilename,
                                              ctPrepFilename=ctPrepFilename,
                                              ntPrepFilename= ntPrepFilename,
                                              arMappingFilename=arMappingFilename,
                                              tryUsingCuda=tryUsingCuda,
                                              useGBSA=useGBSA,
                                              temperature=temperature,
                                              cutoff=cutoff,
                                              useCutoffNonPeriodic=True)
    if verbose:
        sys.stdout.write("Setup time = %.1f Sec\n\n" % (time.time() - timeX))
        sys.stdout.flush()
    timeX=time.time()

    pdbOut=pdbWriter.PDBWriter(amberSystem.pdb)

    if verbose:
        if amberSystem.usingCuda:
            sys.stdout.write("This simulation will be run using the GPU\n")
        else:
            sys.stdout.write("This simulation will be run without using a GPU\n")
        sys.stdout.flush()

    #Create context
    amberContext = amberOpenmmContext.AmberContext(amberSystem, stepSize)


    #Get starting coords
    (simTime, eK0, eP0,
     coordList, velList, forceList) = amberContext.getState(getPositions=True,
                                                            getVelocities=False,
                                                            getForces=True)

    maxTimeLength=len('%.3f' % (stepSize*numSteps))+1
    remark = "%*.3f psec: EK = %9.3f kj   EP = %9.3f kj\n" \
            % (maxTimeLength, simTime, eK0, eP0)
    pdbOut.appendToPdb(pdbOutFilename, coordList,
                       addRemarks='REMARK  %s' % remark,
                       useNM=True)
    if verbose:
        sys.stdout.write(remark)
        sys.stdout.flush()

    #Run simulation and write coords to pdb file
    step = 0
    while step < numSteps:
        #Do stepsPerReport steps
        amberContext.step(stepsPerReport)
        step+=stepsPerReport
        (simTime, eK, eP,
         coordList, velList, forceList) = amberContext.getState(
                                                           getPositions=True,
                                                           getVelocities=False,
                                                           getForces=True)
        remark = "%*.3f psec: EK = %9.3f kj   EP = %9.3f kj\n" \
                % (maxTimeLength, simTime, eK, eP)
        pdbOut.appendToPdb(pdbOutFilename, coordList,
                           addRemarks='REMARK  %s' % remark,
                           useNM=True)
        if verbose:
            sys.stdout.write(remark)
            sys.stdout.flush()
    pdbOut.close(pdbOutFilename)
    return time.time()-time0


def parseCommandLine():
    import getopt
    shortOpts = 'hd:r:a:o:c:n:o:m:'
    longOpts = ['help=', 'pdbIn=', 'param=', 'prep=',
                'ctPrep=', 'ntPrep=', '--arMapping=',
                'pdbOut=']
    opts, args_proper = getopt.getopt(sys.argv[1:], shortOpts, longOpts)
    pdbInFilename = None
    paramFilename = None
    prepFilename = None
    ctPrepFilename = None
    ntPrepFilename = None
    arMappingFilename = None
    pdbOutFilename = None
    for option, parameter in opts:
        if option=='-h': usageError()
        if option=='-d' or option=='--pdbIn': pdbInFilename = parameter
        if option=='-a' or option=='--param': paramFilename = parameter
        if option=='-r' or option=='--prep': prepFilename = parameter
        if option=='-c' or option=='--ctPrep': ctPrepFilename = parameter
        if option=='-n' or option=='--ntPrep': ntPrepFilename = parameter
        if option=='-m' or option=='--arMapping': arMappingFilename = parameter
        if option=='-o' or option=='--pdbOut': pdbOutFilename = parameter
    return (args_proper,
            pdbInFilename,
            paramFilename,
            prepFilename,
            ctPrepFilename,
            ntPrepFilename,
            arMappingFilename,
            pdbOutFilename)

def usageError():
    s = 'usage: %s' % os.path.basename(sys.argv[0])
    sBlanks = len(s)*' '
    print '%s  --pdbIn=PDB_IN_FILENAME            \\' % s
    print '%s  --paramIn=PARM_IN_FILENAME         \\' % sBlanks
    print '%s  --prepIn=PREP_IN_FILENAME          \\' % sBlanks
    print '%s [--ctPrepIn=CTERM_PREP_IN_FILENAME] \\' % sBlanks
    print '%s [--ntPrepIn=NTERM_PREP_IN_FILENAME] \\' % sBlanks
    print '%s [--arMapping=ATOM_RES_MAPPING]      \\' % sBlanks
    print '%s  --pdbOut=PDB_OUT_FILENAME          \\' % sBlanks
    print '%s [--help]' % sBlanks
    sys.exit(1)


if __name__=='__main__':
    dTime=main(verbose=True)
    sys.stdout.write("Total run time = %.1f Sec\n" % dTime)