#!/usr/local/bin/env python

#=============================================================================================
# MODULE DOCSTRING
#=============================================================================================

"""
Illustrate pressure regulation Bus Error.

DESCRIPTION

COPYRIGHT

@author John D. Chodera <jchodera@gmail.com>

All code in this repository is released under the GNU General Public License.

This program is free software: you can redistribute it and/or modify it under
the terms of the GNU General Public License as published by the Free Software
Foundation, either version 3 of the License, or (at your option) any later
version.

This program is distributed in the hope that it will be useful, but WITHOUT ANY
WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A
PARTICULAR PURPOSE.  See the GNU General Public License for more details.
 
You should have received a copy of the GNU General Public License along with
this program.  If not, see <http://www.gnu.org/licenses/>.

TODO

"""

#=============================================================================================
# GLOBAL IMPORTS
#=============================================================================================

import os
import os.path
import sys
import math
import copy
import time

import numpy

import simtk
import simtk.unit as units
import simtk.openmm as openmm

#=============================================================================================
# MAIN AND TESTS
#=============================================================================================

if __name__ == "__main__":
    verbose = True

    # PARAMETERS
    complex_prmtop_filename = 'system.prmtop'
    complex_crd_filename = 'system.crd'

    pressure = 1.0 * units.atmosphere
    temperature = 300.0 * units.kelvin
    collision_rate = 1.0 / units.picosecond
    timestep = 2.0 * units.femtosecond
    barostat_frequency = 10
    nsteps = 50000
    platform_name = 'OpenCL'

    # Load standard systems.
    import simtk.pyopenmm.amber.amber_file_parser as amber
    cutoff = 9.0 * units.angstrom
    reference_system = amber.readAmberSystem(complex_prmtop_filename, mm=openmm, gbmodel=None, nonbondedCutoff=cutoff, nonbondedMethod='CutoffPeriodic', shake='h-bonds')
    [coordinates, box_vectors] = amber.readAmberCoordinates(complex_crd_filename, read_box=True, asNumpy=True)

    # Add barostat.
    barostat = openmm.MonteCarloBarostat(pressure, temperature, barostat_frequency)
    reference_system.addForce(barostat)

    # Get platform.
    platform = openmm.Platform.getPlatformByName(platform_name)

    # Create integrator and context.
    integrator = openmm.LangevinIntegrator(temperature, collision_rate, timestep)
    context = openmm.Context(reference_system, integrator, platform)

    # Set coordinates.
    context.setPositions(coordinates)

    # Set box vectors.
    context.setPeriodicBoxVectors(box_vectors[0], box_vectors[1], box_vectors[2])

    # Minimize.
    print "Minimizing energy..."
    openmm.LocalEnergyMinimizer.minimize(context)

    # Run dynamics.
    print "Running %d steps of dynamics..." % nsteps
    integrator.step(nsteps)

    # Done.
    print "Done."
    del context
    del integrator