#!/usr/local/bin/env python

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

"""
Set up kinase simulation with OpenMM app.

"""

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

import sys
import math
import doctest
import time
import numpy

import simtk.unit as units
import simtk.openmm as openmm
import simtk.pyopenmm.extras.testsystems as testsystems

import netCDF4 as netcdf 

#=============================================================================================
# CONSTANTS
#=============================================================================================

kB = units.BOLTZMANN_CONSTANT_kB * units.AVOGADRO_CONSTANT_NA

#=============================================================================================
# PARAMETERS
#=============================================================================================

timestep = 2.0 * units.femtoseconds
temperature = 298.0 * units.kelvin
pressure = 1.0 * units.atmosphere # pressure for equilibration
gamma = 91.0 / units.picosecond # collision rate
ghmc_nsteps = 10000 # number of steps to generate new uncorrelated sample with GHMC
ghmc_timestep = 0.50 * units.femtoseconds
nsamples = 2000 # number of samples to generate
nequil = 10 # number of NPT equilibration iterations

verbose = True

kT = kB * temperature # thermal energy
beta = 1.0 / kT # inverse temperature

#=============================================================================================
# Create system to simulate.
#=============================================================================================

import simtk.openmm.app as app

# Load structure from PDB file.
if verbose: print "Loading structure from PDB file..."
#pdb = app.PDBFile('src_tbosutinib_c4+d2_refine_waters_tls_38.pdb')
#pdb = app.PDBFile('part1.pdb')
#pdb = app.PDBFile('src-modelled.pdb')
#pdb = app.PDBFile('non-protein-components.pdb')
pdb = app.PDBFile('src-modelled-with-waters.pdb')

# Load forcefield for protein and solvent.
# AMBER 99SB-ILDN and TIP3P water
if verbose: print "Loading forcefield definition..."
forcefield = app.ForceField('amber99sbildn.xml', 'tip3p.xml')

# Create new Modeller instance.
if verbose: print "Creating new Modeller instance..."
modeller = app.Modeller(pdb.topology, pdb.positions)

# Add protons.
if verbose: print "Protonating protein..."
pH = 7.0 
modeller.addHydrogens(forcefield, pH=pH)

# Add solvent to specified box dimensions.
if verbose: print "Adding solvent..."
modeller.addSolvent(forcefield, model='tip3p', padding=1.0*units.nanometers)

# Get new topology and coordinates.
if verbose: print "Generating new topology and coordinates..."
topology = modeller.getTopology()
positions = modeller.getPositions()

# Convert positions to numpy.
positions = units.Quantity(numpy.array(positions / positions.unit), positions.unit)

# Create OpenMM System.
if verbose: print "Creating OpenMM system..."
nonbondedMethod = app.CutoffPeriodic
cutoff = 9.0 * units.angstroms
system = forcefield.createSystem(topology, nonbondedMethod=nonbondedMethod, nonbondedCutoff=cutoff, constraints=app.HBonds, rigidWater=True, removeCMMotion=False)

# 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('initial.pdb', 'w'))