#!/bin/env python

import os, sys
import time

import simtk.unit as unit
import simtk.pyopenmm.amber.zander as zander
import simtk.openmm as openmm
try:
    openMM_version=openmm.Platform.getOpenMMVersion()
except AttributeError:
    raise Exception("must use OpenMM version 2.0 or better")

scriptName=os.path.splitext(os.path.split(sys.argv[0])[-1])[0]
mdin = None
mdout = "protein_%s.out" % scriptName
prmtop = "protein.top"
inpcrd = "protein_minimized.rst"
restrt = "protein_%s.rst" % scriptName
mdcrd = "protein_%s.trj" % scriptName
pdbout = "protein_%s_trj.pdb" % scriptName

overrideParams={}
overrideParams['totalSteps'] = 1000000
overrideParams['stepsPerReport'] = 1000
overrideParams['stepsPerCoordDump'] = 0
overrideParams['usePME'] = False
overrideParams['usePBC'] = False
overrideParams['shakeBondsWithH'] = False
overrideParams['shakeAllBonds'] = False
overrideParams['stepSize'] = 1e-3 * unit.picosecond
overrideParams['useLangevin'] = True
overrideParams['langevinTemperature'] = 300.0 * unit.kelvin
overrideParams['frictionCoeff'] = 5 / unit.picosecond
overrideParams['nonbondedCutoff'] = 20.0 * unit.angstrom
overrideParams['cMMotionRemoverFrequency'] = 100
overrideParams['preferredPlatformName'] = "OpenCL"
overrideParams['preferredGPU'] = None

import membraneGB_Rev0 as membraneGB
extraCustomGBForce=membraneGB.buildCustomMembraneGBForce
time0 = time.time()
zSim=zander.Zander(mdin, mdout, prmtop, inpcrd, restrt, mdcrd, pdbout,
                   extraCustomGBForce=extraCustomGBForce,
                   overrideParams=overrideParams)

time1 = time.time()
zSim.runSim()

zSim.writeMdout("OpenMM Version %s was used\n" % openMM_version)
time2 = time.time()
zSim.writeMdout("Setup wall clock time: %.1f Sec\n" % ((time1-time0)))
zSim.writeMdout("Non-setup wall clock time: %.1f Sec\n" % ((time2-time1)))
zSim.writeMdout("Total wall clock time: %.1f Sec\n" % ((time2-time0)))

zSim.shutdown()