import moldesign as mdt
import moldesign.units as u

"""
Explanation: <span style="float:right"><a href="http://moldesign.bionano.autodesk.com/" target="_blank" title="About">About</a>&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;<a href="https://github.com/autodesk/molecular-design-toolkit/issues" target="_blank" title="Issues">Issues</a>&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;<a href="http://bionano.autodesk.com/MolecularDesignToolkit/explore.html" target="_blank" title="Tutorials">Tutorials</a>&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;<a href="http://autodesk.github.io/molecular-design-toolkit/" target="_blank" title="Documentation">Documentation</a></span>
</span>

<br>
<center><h1>Example 4: The Dynamics of HIV Protease bound to a small molecule </h1> </center>
This notebook prepares a co-crystallized protein / small molecule ligand structure from the PDB database and prepares it for molecular dynamics simulation. 

Author: Aaron Virshup, Autodesk Research<br>
Created on: August 9, 2016
Tags: HIV Protease, small molecule, ligand, drug, PDB, MD
End of explanation
"""


protease = mdt.from_pdb('3AID')
protease

protease.draw()

"""
Explanation: Contents


I. The crystal structure
A. Download and visualize
B. Try assigning a forcefield


II. Parameterizing a small molecule
A. Isolate the ligand
B. Assign bond orders and hydrogens
C. Generate forcefield parameters


III. Prepping the protein
A. Strip waters
B. Histidine


IV. Prep for dynamics
A. Assign the forcefield
B. Attach and configure simulation methods
D. Equilibrate the protein



I. The crystal structure
First, we'll download and investigate the 3AID crystal structure.
A. Download and visualize
End of explanation
"""


amber_ff = mdt.forcefields.DefaultAmber()
newmol = amber_ff.create_prepped_molecule(protease)

"""
Explanation: B. Try assigning a forcefield
This structure is not ready for MD - this command will raise a ParameterizationError Exception. After running this calculation, click on the Errors/Warnings tab to see why.
End of explanation
"""


sel = mdt.widgets.ResidueSelector(protease)
sel

drugres = mdt.Molecule(sel.selected_residues[0])
drugres.draw2d(width=700, show_hydrogens=True)

"""
Explanation: You should see 3 errors: 
 1. The residue name ARQ not recognized
 1. Atom HD1 in residue HIS69, chain A was not recognized
 1. Atom HD1 in residue HIS69, chain B was not recognized
(There's also a warning about bond distances, but these can be generally be fixed with an energy minimization before running dynamics)
We'll start by tackling the small molecule "ARQ".
II. Parameterizing a small molecule
We'll use the GAFF (generalized Amber force field) to create force field parameters for the small ligand.
A. Isolate the ligand
Click on the ligand to select it, then we'll use that selection to create a new molecule.
End of explanation
"""


drugmol = mdt.tools.set_hybridization_and_saturate(drugres)
drugmol.draw(width=500)

drugmol.draw2d(width=700, show_hydrogens=True)

"""
Explanation: B. Assign bond orders and hydrogens
A PDB file provides only limited information; they often don't provide indicate bond orders, hydrogen locations, or formal charges. These can be added, however, with the add_missing_pdb_data tool:
End of explanation
"""


drug_parameters = mdt.create_ff_parameters(drugmol, charges='gasteiger')

"""
Explanation: C. Generate forcefield parameters
We'll next generate forcefield parameters using this ready-to-simulate structure.
NOTE: for computational speed, we use the gasteiger charge model. This is not advisable for production work! am1-bcc or esp are far likelier to produce sensible results.
End of explanation
"""


dehydrated = mdt.Molecule([atom for atom in protease.atoms if atom.residue.type != 'water'])

"""
Explanation: III. Prepping the protein
Section II. dealt with getting forcefield parameters for an unknown small molecule. Next, we'll prep the other part of the structure.
A. Strip waters
Waters in crystal structures are usually stripped from a simulation as artifacts of the crystallization process. Here, we'll remove the waters from the protein structure.
End of explanation
"""


mdt.guess_histidine_states(dehydrated)

"""
Explanation: B. Histidine
Histidine is notoriously tricky, because it exists in no less than three different protonation states at biological pH (7.4) - the "delta-protonated" form, referred to with residue name HID; the "epsilon-protonated" form aka HIE; and the doubly-protonated form HIP, which has a +1 charge. Unfortunately, crystallography isn't usually able to resolve the difference between these three.
Luckily, these histidines are pretty far from the ligand binding site, so their protonation is unlikely to affect the dynamics. We'll therefore use the guess_histidine_states function to assign a reasonable starting guess.
End of explanation
"""


amber_ff = mdt.forcefields.DefaultAmber()
amber_ff.add_ff(drug_parameters)
sim_mol = amber_ff.create_prepped_molecule(dehydrated)

"""
Explanation: IV. Prep for dynamics
With these problems fixed, we can succesfully assigne a forcefield and set up the simulation.
A. Assign the forcefield
Now that we have parameters for the drug and have dealt with histidine, the forcefield assignment will succeed:
End of explanation
"""


sim_mol.set_energy_model(mdt.models.OpenMMPotential, implicit_solvent='obc', cutoff=8.0*u.angstrom)
sim_mol.set_integrator(mdt.integrators.OpenMMLangevin, timestep=2.0*u.fs)
sim_mol.configure_methods()

"""
Explanation: B. Attach and configure simulation methods
Armed with the forcefield parameters, we can connect an energy model to compute energies and forces, and an integrator to create trajectories:
End of explanation
"""


mintraj = sim_mol.minimize()
mintraj.draw()

traj = sim_mol.run(20*u.ps)

viewer = traj.draw(display=True)
viewer.autostyle()

"""
Explanation: C. Equilibrate the protein
The next series of cells first minimize the crystal structure to remove clashes, then heats the system to 300K.
End of explanation
"""


