#!/usr/bin/env python
# coding: utf-8

# # Simulate an Interchange with Amber
# 
# <details>
#     <summary><small>▼ Click here for dependency installation instructions</small></summary>
#     The simplest way to install dependencies is to use the Interchange examples environment. From the root of the cloned openff-interchange repository:
#     
#     conda env create --name interchange-examples --file devtools/conda-envs/examples_env.yaml 
#     conda activate interchange-examples
#     pip install -e .
#     cd examples/amber
#     jupyter notebook amber.ipynb
#     
# </details>
# 
# In this example, we'll quickly construct an `Interchange` and then run a simulation in Amber. 

# We need an `Interchange` to get started, so let's put that together quickly. For more explanation on this process, take a look at the [packed_box] and [protein_ligand] examples.
# 
# [packed_box]: https://github.com/openforcefield/openff-interchange/tree/main/examples/packed_box
# [protein_ligand]: https://github.com/openforcefield/openff-interchange/tree/main/examples/protein_ligand

# In[ ]:


import mdtraj
import nglview
import openmm.app
from openff.toolkit import ForceField, Molecule, Topology
from openff.toolkit.utils import get_data_file_path

from openff.interchange import Interchange

# Read a structure from the Toolkit's test suite into a Topology
pdbfile = openmm.app.PDBFile(get_data_file_path("systems/packmol_boxes/propane_methane_butanol_0.2_0.3_0.5.pdb"))
molecules = [Molecule.from_smiles(smi) for smi in ["CCC", "C", "CCCCO"]]
off_topology = Topology.from_openmm(pdbfile.topology, unique_molecules=molecules)

# Construct the Interchange with the OpenFF "Sage" force field
interchange = Interchange.from_smirnoff(
    force_field=ForceField("openff-2.0.0.offxml"),
    topology=off_topology,
)
interchange.positions = pdbfile.positions


# Tada! A beautiful solvent system:

# In[ ]:


interchange.visualize("nglview")


# 
# ## Run a simulation
# 
# We need Amber input files to run our simulation. `Interchange.to_amber` takes a (string) prefix as an argument and wraps three other methods that each write out a file needed for running a simulation in Amber:
# * `mysim.prmtop` stores the chemical topology and physics paramaters
# * `mysim.inpcrd` file stores coordinates
# * `mysim_pointenergy.in` tells `sander` how a single-point energy "simulation" should be run

# In[ ]:


interchange.to_amber("mysim")

get_ipython().system('ls mysim*')


# To get a proper simulation with a trajectory, we'll also need an input file to describe the simulation parameters a a few other details, like a thermostat and what information to write to files:

# In[ ]:


amber_in = """Basic Amber control file
&cntrl
  imin=0,                ! Run molecular dynamics.
  ntx=1,                 ! Take positions from input and generate velocities
  nstlim=500,            ! Number of MD-steps to be performed.
  dt=0.001,              ! Time step (ps), use a low 1 ps timestep to be safe
  tempi=300.0,           ! Initial temperature for velocity generation
  temp0=300.0,           ! Thermostat temperature
  cut=9.0,               ! vdW cutoff (Å)
  fswitch=8.0            ! vdW switching function start point (Å)
  igb=0,                 ! Don't use a Generalized Born model
  ntt=3, gamma_ln=2.0,   ! Temperature scaling using Langevin dynamics with the collision frequency in gamma_ln (1/ps)
  ntp=0,                 ! No pressure scaling
  iwrap=1,               ! Wrap trajectory coordinates to stay in box
  ioutfm=1,              ! Write out netcdf trajectory
  ntwx=10,               ! Frequency to write coordinates
  ntpr=50,               ! Frequency to log energy info
  ntc=2,                 ! Constraints on Hydrogen-involving bonds
/
"""
with open("amber.in", "w") as f:
    f.write(amber_in)


# Run the simulation with Sander:

# In[ ]:


get_ipython().system('sander                      -O                       -i amber.in              -p mysim.prmtop          -c mysim.inpcrd          -x trajectory.nc')


# And finally we can visualize!

# In[ ]:


traj = mdtraj.load("trajectory.nc", top=mdtraj.load_prmtop("mysim.prmtop"))
nglview.show_mdtraj(traj)