#!/usr/local/bin/env python

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

"""
Set up Kate Engel's EGFR structure for simulation with OpenMM/

"""

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

import sys
import math
import doctest
import time
import numpy

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

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

# Load AMBER 99SB-ILDN forcefield and AMBER 99 OBC GB model.
print "Reading forcefields..."
forcefield = app.ForceField('amber99sbildn.xml', 'amber99_obc.xml')

# Load modelled PDB file.
print "Reading PDB file..."
pdbfile = app.PDBFile('highresegfraligned-modelled.pdb')

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

# Add hydrogens.
print "Adding hydrogens..."
modeller.addHydrogens(forcefield)

# Create system.
print "Creating OpenMM System object..."
system = forcefield.createSystem(modeller.topology, implicitSolvent=app.OBC2, constraints=app.HBonds)

# Create 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(modeller.topology, system, integrator)

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

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

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

# Run dynamics
nsteps = 1000
simulation.reporters.append(app.StateDataReporter(sys.stdout, 10, step=True, potentialEnergy=True, temperature=True))
simulation.step(nsteps)

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