#!/usr/bin/env python

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


#Simulation options
numParticles=3
stepSize=0.020      #ps
totalTime=500.0     #ps
stepsPerReport=1000
#randomSeed=None
randomSeed=1
verbose=True
tryUsingGpu=True

#Parameters
mass_Ar     = 39.9         #AMU
q_Ar        = 0.0          #q
sigma_Ar    = 3.350 / 10   #nm
epsilon_Ar  = 0.001603     #kJ/mol


import  os, sys, math, time

import simtk.utils.manageGPUs as gMan
(mm, usingCuda) = gMan.importBestOpenMMLib(tryUsingCuda=tryUsingGpu)

def appendCoordsToPDB(fOut, step, eK, eP, coords,
                      modelFrameNumber):
    """Append one coord set to the open PDB file"""
    fOut.write("REMARK  step %6d: EK = %9.3f j   EP = %9.3f j\n"
                % (step, 1000*eK, 1000*eP))
    fOut.write("MODEL     %d\n" % modelFrameNumber)
    for ii in range(len(coords)):
        fOut.write("ATOM   %4d  AR  ARG  %4d    %8.3f%8.3f%8.3f\n"
                   % (ii, ii,
                      10*coords[ii][0],
                      10*coords[ii][1],
                      10*coords[ii][2]) )
    fOut.write("ENDMDL\n")


def main():
    """Main loop, builds OpenMM system, adds particles and parameters,
       runs and write coords at PDB file"""
    time0=time.time()
    try:
        outFilename=sys.argv[1]
    except IndexError:
        usageError()

    #Open out PDB file
    fOut=open(outFilename, 'w')

    if 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 system
    system=mm.System()

    #Add particles to system
    initPositions=[]
    nb = mm.NonbondedForce()
    for ii in range(numParticles):
        system.addParticle(mass_Ar)
        nb.addParticle(q_Ar, sigma_Ar, epsilon_Ar)
        initPositions.append( (0.5*ii, 0, 0) )
    system.addForce(nb)
    nb.thisown=False

    #Select simple integrator
    integrator = mm.VerletIntegrator(stepSize)
    #Make a context for this system
    context=mm.OpenMMContext(system, integrator)
    #Staring config
    context.setPositions(initPositions)

    if verbose:
        sys.stdout.write("Run to %.3f ps\n" % totalTime)
    #Setup simulation, and write first coord set
    step=0
    (simTime, eK0, eP0,
     coords, velList, forceList) = context.getState(getPositions=True,
                                                    getVelocities=False,
                                                    getForces=False)
    if verbose:
        sys.stdout.write("%8.2f ps: Potential Energy = %8.3f j\n" % (stepSize*step, 1000*eP0))
        sys.stdout.flush()

    appendCoordsToPDB(fOut, step, eK0, eP0, coords, 0)

    #Run simulation, and write coords
    count=0
    while round(context.getTime(), 3) < round(totalTime, 3):
        #Do stepsPerReport steps
        integrator.step(stepsPerReport)
        step+=stepsPerReport
        (simTime, eK, eP,
         coords, velList, forceList) = context.getState(getPositions=True,
                                                        getVelocities=False,
                                                        getForces=False)
        if verbose:
            sys.stdout.write("%8.2f ps: Potential Energy = %8.3f j\n" % (stepSize*step, 1000*eP))
            sys.stdout.flush()
        appendCoordsToPDB(fOut, step, eK, eP, coords, count+1)
        count+=1
    fOut.write("END\n")
    fOut.close()

    return time.time()-time0


def usageError():
    sys.stdout.write('usage: %s outFilename\n'
                     % os.path.basename(sys.argv[0]))
    sys.exit(1)

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