# coding: utf-8

# In[ ]:

from __future__ import print_function
get_ipython().magic(u'matplotlib inline')
import matplotlib.pyplot as plt
import itertools
import mdtraj as md
import mdtraj.testing


# Load up some example data. This is a little 28 residue peptide

# In[ ]:

t = md.load_pdb('http://www.rcsb.org/pdb/files/2EQQ.pdb')
print(t)


# `md.baker_hubbard` idenfies hydrogen bonds baced on cutoffs
# for the Donor-H...Acceptor distance and angle. The criterion employed
# is $\theta > 120$ and $r_\text{H...Acceptor} < 2.5 A$ in
# at least 10% of the trajectory. The return value is a list of the 
# indices of the atoms (donor, h, acceptor) that satisfy this criteria.

# In[ ]:

hbonds = md.baker_hubbard(t, periodic=False)
label = lambda hbond : '%s -- %s' % (t.topology.atom(hbond[0]), t.topology.atom(hbond[2]))
for hbond in hbonds:
    print(label(hbond))


# Let's compute the actual distances between the donors and acceptors

# In[ ]:

da_distances = md.compute_distances(t, hbonds[:, [0,2]], periodic=False)


# And plot a histogram for a few of them

# In[ ]:

color = itertools.cycle(['r', 'b', 'gold'])
for i in [2, 3, 4]:
    plt.hist(da_distances[:, i], color=next(color), label=label(hbonds[i]), alpha=0.5)
plt.legend()
plt.ylabel('Freq');
plt.xlabel('Donor-acceptor distance [nm]')


# In[ ]: