#!/usr/local/bin/env python

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

"""
Run molecular dynamics simulation of Kate Engel's EGFR docked system in explicit solvent using OpenMM.

"""

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

import sys
import math
import doctest
import time
import numpy

#=============================================================================================
# MAIN
#=============================================================================================

# Define AMBER filenames.
prmtop_filename = 'complex.prmtop'
inpcrd_filename = 'complex.crd'
pdb_filename = 'complex.pdb'

# Import OpenMM, units package, and application layer.
import simtk.unit as units
import simtk.openmm as openmm
import simtk.openmm.app as app

# Load the AMBER system.                                                                                                                       
print "Creating AMBER system..."                                                                                                               
inpcrd = app.AmberInpcrdFile(inpcrd_filename)                                                                                                  
prmtop = app.AmberPrmtopFile(prmtop_filename)                                                                                                  
system = prmtop.createSystem(nonbondedMethod=app.CutoffPeriodic, constraints=app.HBonds)                                   
positions = inpcrd.getPositions()

# Create new OpenMM Modeller instance.
print "Creating OpenMM Modeller instance..."
pdbfile = app.PDBFile(pdb_filename)
modeller = app.Modeller(prmtop.topology, pdbfile.positions)

# Create Langevin integrator.
print "Creating integrator..."
temperature = 298.0 * units.kelvin
collision_rate = 9.1 / units.picosecond
timestep = 2.0 * units.femtoseconds
integrator = openmm.LangevinIntegrator(temperature, collision_rate, timestep)

# Create OpenMM app Simulation object.
print "Creating Simulation object..."
simulation = app.Simulation(prmtop.topology, system, integrator)

# Set positions.
print "Setting positions..."
simulation.context.setPositions(positions)

# Minimize energy.
print "Minimizing energy..."
simulation.minimizeEnergy(maxIterations=100)

# Write PDB file.
print "Writing PDB..."
positions = simulation.context.getState(getPositions=True).getPositions()
app.PDBFile.writeFile(prmtop.topology, positions, open('minimized.pdb', 'w'))

nsteps_report = 500
nsteps = 10 * nsteps_report

# Add PDB reporter to write a frame every 1 ps.
simulation.reporters.append(app.PDBReporter('md.pdb', nsteps_report)) 

# Run 1 ns of dynamics.
simulation.reporters.append(app.StateDataReporter(sys.stdout, 50, step=True, potentialEnergy=True, temperature=True)) 
simulation.step(nsteps)