#!/usr/bin/env python

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


#Simulation options
stepSize=0.001      #ps
totalSteps=2
stepsPerReport=1
#randomSeed=None
randomSeed=1
verbose=True
tryToUseCuda=True

import  os, sys, math


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


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

    (args_proper,
     pdbFilename,
     paramFilename,
     prepFilename,
     ctPrepFilename,
     ntPrepFilename) = parseCommandLine()
    if not pdbFilename or not paramFilename or not prepFilename:
         usageError()

    amberSystem=amberOpenmmSystem.AmberSystem(pdbFilename,
                                              paramFilename,
                                              prepFilename,
                                              ctPrepFilename=ctPrepFilename,
                                              ntPrepFilename= ntPrepFilename,
                                              tryToUseCuda=tryToUseCuda,
                                              temperature=0.0)

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

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


    #Setup simulation, and write first coord set
    (simTime, eK0, eP0,
     coordList, velList, forceList) = amberContext.getState(getPositions=True,
                                                            getVelocities=False,
                                                            getForces=True)
    sys.stdout.write("ENERG %f %f\n" % (eK0, eP0))
    for ii in range(len(coordList)):
        sys.stdout.write("COORD %d %f %f %f\n" % (ii,
                                                  coordList[ii][0], 
                                                  coordList[ii][1], 
                                                  coordList[ii][2]))
    for ii in range(len(coordList)):
        sys.stdout.write("FORCE %d %f %f %f\n" % (ii,
                                                  forceList[ii][0], 
                                                  forceList[ii][1], 
                                                  forceList[ii][2]))
    sys.stdout.flush()

    #Run simulation, and write coords
    step = 0
    while step <= totalSteps:
        #Do stepsPerReport steps
        amberContext.step(stepsPerReport)
        step+=stepsPerReport
        (simTime, eK, eP,
         coordList, velList, forceList) = amberContext.getState(
                                                           getPositions=True,
                                                           getVelocities=False,
                                                           getForces=True)
        sys.stdout.write("ENERG %f %f\n" % (eK, eP))
        for ii in range(len(coordList)):
            sys.stdout.write("COORD %d %f %f %f\n" % (ii,
                                                      coordList[ii][0], 
                                                      coordList[ii][1], 
                                                      coordList[ii][2]))
        for ii in range(len(coordList)):
            sys.stdout.write("FORCE %d %f %f %f\n" % (ii,
                                                      forceList[ii][0], 
                                                      forceList[ii][1], 
                                                      forceList[ii][2]))
        sys.stdout.flush()

    return time.time()-time0


def parseCommandLine():
    import getopt
    shortOpts = 'hd:r:a:o:c:n:'
    longOpts = ['help=', 'pdbIn=', 'paramIn=', 'prepIn=', 'ctPrepIn=', 'ntPrepIn=']
    opts, args_proper = getopt.getopt(sys.argv[1:], shortOpts, longOpts)
    pdbFilename = None
    paramFilename = None
    prepFilename = None
    ctPrepFilename = None
    ntPrepFilename = None
    for option, parameter in opts:
        if option=='-h': usageError()
        if option=='-d' or option=='--pdbIn': pdbFilename = parameter
        if option=='-a' or option=='--paramIn': paramFilename = parameter
        if option=='-r' or option=='--prepIn': prepFilename = parameter
        if option=='-c' or option=='--ctPrepIn': ctPrepFilename = parameter
        if option=='-n' or option=='--ntPrepIn': ntPrepFilename = parameter
    return (args_proper,
            pdbFilename,
            paramFilename,
            prepFilename,
            ctPrepFilename,
            ntPrepFilename)

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  [--help]' % sBlanks
    sys.exit(1)


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