from __future__ import print_function

from simtk.openmm import app
import simtk.openmm as mm
from simtk import unit
import sys
import mbpol

pdb = app.PDBFile("water14_cluster.pdb")

forcefield = app.ForceField("mbpol.xml")
# use tip4p
#forcefield = app.ForceField("tip4pfb.xml")

system = forcefield.createSystem(pdb.topology, nonbondedMethod=app.CutoffNonPeriodic, nonBondedCutoff=1e3*unit.nanometer)
integrator = mm.VerletIntegrator(0.00001*unit.femtoseconds)

platform = mm.Platform.getPlatformByName('Reference')
simulation = app.Simulation(pdb.topology, system, integrator, platform)
simulation.context.setPositions(pdb.positions)
simulation.context.computeVirtualSites()

state = simulation.context.getState(getForces=True, getEnergy=True)
potential_energy = state.getPotentialEnergy()
potential_energy.in_units_of(unit.kilocalorie_per_mole)

kilocalorie_per_mole_per_angstrom = unit.kilocalorie_per_mole/unit.angstrom
for f in state.getForces():
    print(f.in_units_of(kilocalorie_per_mole_per_angstrom))

from simtk.openmm import LocalEnergyMinimizer

LocalEnergyMinimizer.minimize(simulation.context, 1e-1)

state = simulation.context.getState(getForces=True, getEnergy=True, getPositions=True)
potential_energy = state.getPotentialEnergy()
potential_energy.in_units_of(unit.kilocalorie_per_mole)

kilocalorie_per_mole_per_angstrom = unit.kilocalorie_per_mole/unit.angstrom
for f in state.getForces():
    print(f.in_units_of(kilocalorie_per_mole_per_angstrom))

state.getPositions()

simulation.context.setVelocitiesToTemperature(300*unit.kelvin)
# Equilibrate
simulation.step(10)

simulation.reporters.append(app.StateDataReporter(sys.stdout, 10, step=True, 
    potentialEnergy=True, temperature=True, progress=True, remainingTime=True, 
    speed=True, totalSteps=110, separator='\t'))

simulation.reporters.append(app.PDBReporter('trajectory.pdb', 20))

simulation.step(100)

!head trajectory.pdb

!echo Number of lines: `wc -l trajectory.pdb`