#!/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 TIP3P water model.
print "Reading forcefields..."
forcefield = app.ForceField('amber99sbildn.xml', 'tip3p.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)

# Add TIP3P solvent.
print "Adding solvent..."
modeller.addSolvent(forcefield, model='tip3p', padding=1.0*units.nanometer)

# Create system.
print "Creating OpenMM System object..."
system = forcefield.createSystem(modeller.topology, nonbondedMethod=app.PME, 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=100)

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