# This cell sets up both the python and notebook environments
%matplotlib inline  
import moldesign as mdt  # import the buckyball package
from moldesign import units as u  # import the buckyball unit system

"""
Explanation: <a href="http://moldesign.bionano.autodesk.com" target="_blank"><img src="img/Top.png"></a>
<center><h1>MDT Quickstart<br> or,
A Hitchhiker's Guide to the Chemical Modeling Universe</h1></center>
This notebook gives you the Molecular Design Toolkit crash course.
Click on the following cell and press Shift + Return to execute it.
End of explanation
"""


my_first_molecule = mdt.from_name('morphine')
my_first_molecule.draw()

"""
Explanation: Molecules
Let's start with something relaxing.
Execute each cell by selecting it and pressing Shift + Return
End of explanation
"""


my_first_molecule

"""
Explanation: You've just created your first molecule in buckyball - the bb.from_name method can translate many common chemical names, as well as IUPAC nomenclature, into 2D and 3D chemical structures.
If morphine is a little intimdating, you can try making caffeine, or even just water, by editing the code in the above cell. You can edit any code in this notebook at any time, and press Shift + Return to execute the new commands.
Let's take a look at your Molecule's attributes. Try using Jupyter's live autocompletion - typing my_first_molecule.[tab], for instance, will show you a list of the object's attributes.
End of explanation
"""


print 'This molecule has %d atoms' % my_first_molecule.num_atoms
for atom in my_first_molecule.atoms[:5]:
    print 'Atom {atom} (atomic number {atom.atnum}) weighs {atom.mass}, and ' \
    'is bonded to {atom.num_bonds} other atoms.'.format(atom=atom)

my_first_molecule.atoms[-1]

"""
Explanation: Every chemical system in buckyball is called a Molecule. These objects store information about atoms, chemical bonds, 3D positions, and chemical properties. A Molecule object is more than just a data store, though - they also have methods for simulating, modeling, and visualizing this information.
A list of the molecule's atoms are stored at my_first_molecule.atoms. Let's print out some custom information about the first five atoms:
End of explanation
"""


mol = mdt.from_smiles('CCCC')
mdt.ui.GeometryBuilder(mol)

"""
Explanation: Manipulating 3D structures
Sometimes, a 3D structure needs to be fixed - perhaps the wrong rotamer was generated, or someone gave you a non-planar geometry of a planar molecule.
For instance, the OpenBabel SMILES converter will generate a gauche-conformation alkane, when we really want an anti conformation. Let's use the GeometryBuilder to fix it - click on the central dihedral bond, and adjust the angle to 180º, either with the slider or by typing "180".
End of explanation
"""


hivprotease = mdt.read('data/3AID.pdb')
hivprotease

"""
Explanation: Readin', Writin', and Introspectin'
Of course, you'll probably want to import and export your own molecular structures. Buckyball can read and write molecules from a variety of 3D formats, including pdb, sdf, mol2, cif, and xyz.
Let's start by reading in a file that was downloaded from the PDB database (you can also use bb.from_pdb(3AID) to download it directly from the PDB site). We can get some quick information about the molecule in the file just by its name on the last line of the cell.
End of explanation
"""


hivprotease.draw()

"""
Explanation: This is a crystal structure of a protein (HIV-1 protease) that's bound to a small organic drug molecule (that's the single unknown residue type). In addition to the protein and drug, there are also several water molecules hanging around the structure.
End of explanation
"""


drugresidue = hivprotease.chains['A'].residues['ARQ401']

drug = mdt.Molecule(drugresidue)
drug.draw3d()

"""
Explanation: Let's pull the drug molecule out of the structure and save it to disk. First, click on one of the atoms in the drug. You'll see that it's Residue ARQ401 in Chain A. We can grab this residue by going into the molecule's secondary structure:
End of explanation
"""


drug.write(filename='/tmp/drugmol.xyz')

!cat /tmp/drugmol.xyz

"""
Explanation: Now, we'll write it to disk:
End of explanation
"""


mol = mdt.from_smiles('C1=CC=CC=C1')
mol.draw()

"""
Explanation: Simulation and modeling
As a quick first example, let's take a look at a quantum mechanical model of benzene.
This time, we'll construct the molecule in yet another way - by using a SMILES string.
End of explanation
"""


mol.set_energy_model(mdt.models.PySCFPotential(theory='rhf',basis='sto-3g'))
mol.calculate()

"""
Explanation: Next, we associate the quantum mechanical model with the molecule, and run a calculation:
End of explanation
"""


mol.draw_orbitals()

"""
Explanation: There's a bunch of data here, which we'll dive into in other notebooks. For now, we can visualize the results of the calculation:
End of explanation
"""


viewer = hivprotease.draw3d()
hdonors = [atom for atom in hivprotease.atoms
           if atom.elem in ('O','N') and atom.residue.type == 'protein']

viewer.add_style(style='vdw', atoms=hdonors, radius=0.5)
viewer.highlight_atoms(atoms=[a for a in hdonors if a.distance(drugresidue) <= 4.0*u.angstrom])

viewer

"""
Explanation: You can visualize two different kinds of orbitals here: the canonical molecular orbitals, often just referred to as "molecular orbitals" (MOs); and atomic orbitals, which were combined to create the MOs.
Visualization API
You can access and create molecule visualizations in the notebook as well. 
First, let's draw some more information into our visualization. We'll draw all the protein's potential hydrogen bond donors (basically Oxygens and Nitrogens) as spheres, and highlight those close to the ligand.
End of explanation
"""


contact_view = bb.viewer.make_contact_view(drugresidue,view_radius=3.5*u.angstrom,
                                    contact_radius=2.25*u.angstrom,
                                          width=600,height=400)
geom_view = bb.viewer.GeometryViewer(hivprotease, width=600,height=300)
geom_view.cartoon(opacity=0.7)
geom_view.licorice(atoms=drugresidue,radius=0.5)
for residue,color in contact_view.colored_residues.iteritems():
    geom_view.licorice(atoms=residue,color=color)

# We use the ipywidgets package to arrange the three UI elements:
import ipywidgets as ipy
viewer_group = ipy.VBox([geom_view,contact_view])
selector = bb.ui.SelectionGroup([ipy.HBox([viewer_group,bb.ui.AtomInspector()])])
selector

"""
Explanation: The cell below constructs a quick visualization of the HIV drug's contacts with the surrounding protein residues, with 2D rendering, a 3D rendering, and an inspector that provides information about atoms that the user clicks on:
End of explanation
"""


