#!/usr/bin/env python
import sys

import simtk.openmm.openmm as mm
import simtk.unit as unit

import dumpOutput
import amber99 as params

stepSize  = 1.0 * unit.femtosecond
totalTime = 1.0 * unit.picosecond
stepsPerReport = 20

fPDB = open('NaCl.pdb', 'w')

system = mm.System()
nb = mm.NonbondedForce()
system.addForce(nb)

atomNames = []
initPositions = [] * unit.angstrom
for ii in range(2):
    atomNames.append('CL')
    system.addParticle(params.mass_Cl)
    nb.addParticle(params.q_Cl, params.sigma_Cl, params.epsilon_Cl)
    initPositions.append((0, 9*ii, 0) * unit.angstrom)

    atomNames.append('NA')
    system.addParticle(params.mass_Na)
    nb.addParticle(params.q_Na, params.sigma_Na, params.epsilon_Na)
    initPositions.append((5, -9*ii, 0) * unit.angstrom)

integrator = mm.VerletIntegrator(stepSize)
context = mm.Context(system, integrator)
context.setPositions(initPositions)

count = 0
simTime = 0 * unit.picosecond
while simTime < totalTime:
    integrator.step(stepsPerReport)
    count += 1
    state = context.getState(getPositions=True,
                             getVelocities=False,
                             getForces=False,
                             getEnergy=True,
                             getParameters=False)
    simTime = state.getTime()
    eK = state.getKineticEnergy()
    eP = state.getPotentialEnergy()
    coords = state.getPositions()
    dumpOutput.writeEnergyAndCoords(fPDB,
                                    stepsPerReport*count,
                                    eK, eP, coords,
                                    simTime, count, atomNames)

dumpOutput.closePDB(fPDB)