traj = pt.Trajectory('step5_production.dcd', '../step3_pbcsetup.xplor.ext.psf')
print(traj)

"""
Explanation: Read in the trajectory file (.dcd) and the topology file (.psf)
End of explanation
"""


goo  = pt.rdf(traj, solvent_mask=':TIP3@OH2', solute_mask=':TIP3@OH2', bin_spacing=0.05, maximum=8.)
goh1 = pt.rdf(traj, solvent_mask=':TIP3@OH2', solute_mask=':TIP3@H1',  bin_spacing=0.05, maximum=8.)
goh2 = pt.rdf(traj, solvent_mask=':TIP3@OH2', solute_mask=':TIP3@H2',  bin_spacing=0.05, maximum=8.)

"""
Explanation: Calculate the water-water radial distribution function. In statistical mechanics, the radial distribution function, (or pair correlation function) g(r) in a system of particles (atoms, molecules, colloids, etc.), describes how density varies as a function of distance from a reference particle.
$$g(r)=\frac{1}{N\rho}\sum_{i=1}^N\sum_{k\neq i}^N\langle\delta(r+r_k-r_i)\rangle$$
where $\rho$ is the number density, $\delta()$ is blah
End of explanation
"""


plt.plot(goo[0],goo[1])
plt.plot(goo[0],goh1[1])
plt.plot((0, 8), (1, 1), 'k--')
plt.xlim([1.5,8])
plt.ylim([0,3])
plt.xlabel('$r_{ox} / \AA$', fontsize=14)
plt.ylabel('$g(r_{ox})$', fontsize=14)
plt.show()

"""
Explanation: Plot the g(r) versus r for O-O and O-H
End of explanation
"""


gon     = pt.rdf(traj, solvent_mask=':TIP3@OH2', solute_mask=':DOP@N23', bin_spacing=0.05, maximum=8.)
goh231  = pt.rdf(traj, solvent_mask=':TIP3@OH2', solute_mask=':DOP@H231', bin_spacing=0.05, maximum=8.)
goo1    = pt.rdf(traj, solvent_mask=':TIP3@OH2', solute_mask=':DOP@O1', bin_spacing=0.05, maximum=8.)
goh4    = pt.rdf(traj, solvent_mask=':TIP3@OH2', solute_mask=':DOP@H4', bin_spacing=0.05, maximum=8.)
ghwn23  = pt.rdf(traj, solvent_mask=':TIP3@H1', solute_mask=':DOP@N23', bin_spacing=0.05, maximum=8.)
ghwh231 = pt.rdf(traj, solvent_mask=':TIP3@H1', solute_mask=':DOP@H231', bin_spacing=0.05, maximum=8.)
ghwo1   = pt.rdf(traj, solvent_mask=':TIP3@H1', solute_mask=':DOP@O1', bin_spacing=0.05, maximum=8.)
ghwh4   = pt.rdf(traj, solvent_mask=':TIP3@H1', solute_mask=':DOP@H4', bin_spacing=0.05, maximum=8.)

"""
Explanation: Calculate molecule - water rdf
End of explanation
"""


plt.plot(gon[0],gon[1],label='OW...NH$_{3}$')
plt.plot(goh231[0],goh231[1],label='OW...HN')
plt.plot(goo1[0],goo1[1],label='OW...O1')
plt.plot(goh4[0],goh4[1],label='OW...HO1')
plt.plot((0, 8), (1, 1), 'k--')
plt.xlim([1.5,8])
plt.ylim([0,0.0015])
plt.xlabel('$r_{ox} / \AA$', fontsize=14)
plt.ylabel('$g(r_{ox})$', fontsize=14)
plt.legend(loc='upper right',frameon=False, bbox_to_anchor=(1.35, 0.95),numpoints=1)
plt.show()

"""
Explanation: Plot the g(r) versus r for Ow-x
End of explanation
"""


plt.plot(ghwn23[0],ghwn23[1],label='HW...NH$_{3}$')
plt.plot(ghwh231[0],ghwh231[1],label='HW...HN')
plt.plot(ghwo1[0],ghwo1[1],label='HW...O1')
plt.plot(ghwh4[0],ghwh4[1],label='HW...HO1')
plt.plot((0, 8), (1, 1), 'k--')
plt.xlim([1.5,8])
plt.ylim([0,0.0008])
plt.xlabel('$r_{hwx} / \AA$', fontsize=14)
plt.ylabel('$g(r_{hwx})$', fontsize=14)
plt.legend(loc='upper right',frameon=False, bbox_to_anchor=(1.35, 0.95),numpoints=1)
plt.show()

"""
Explanation: Plot the g(r) versus r for Hw-x
End of explanation
"""


data = pt.calc_pairdist(traj, mask=':TIP3@OH2', mask2=':DOP@N23')
print(data[:,0][:1])

atom_indices = pt.select(':DOP@N23', traj.top)
print(atom_indices)

"""
Explanation: Print the distance between dop (residue 0) and residue 1 (H2O)
End of explanation
"""


dataset = pt.dihedral(traj, ':DOP@N23 :DOP@C2 :DOP@C3 :DOP@C1')

"""
Explanation: Calculate a dihedral angle as a function of time
End of explanation
"""


time = [i*0.002 for i in range(len(dataset))]
plt.plot(time,dataset, 'ro', label='$\phi$')
plt.xlabel('Time / ns', fontsize=14)
plt.ylabel('Angle / $\circ$', fontsize=14)
plt.legend(loc='upper right',frameon=False, bbox_to_anchor=(1.2, 0.95),numpoints=1)
plt.show()

hist0 = plt.hist(dataset,25)
tmp = plt.setp(hist0[2], 'facecolor', 'r', 'alpha', 0.75)

"""
Explanation: Plot the results
End of explanation
"""


