import sys; sys.path.append('.')
from plotdialog import PlotDialog
del sys.path[-1]

class RamachandranPlot(PlotDialog):

	def __init__(self, m):
		PlotDialog.__init__(self)
		self.mol = m
		self.phi = []
		self.psi = []
		self.residues = []
		for r in m.residues:
			if r.phi is None or r.psi is None:
				continue
			self.phi.append(r.phi)
			self.psi.append(r.psi)
			self.residues.append(r)
		self.selectedIndices = []
		ax = self.add_subplot(1,1,1)
		self.subplot = ax
		self._displayData()
		self.registerPickHandler(self.onPick)

		from chimera import triggers
		self.selHandler = triggers.addHandler("selection changed",
							self._selectionChangedCB,
							None)

	def _displayData(self):
		colors = ['b'] * len(self.phi)
		for n in self.selectedIndices:
			colors[n] = 'r'
		ax = self.subplot
		ax.clear()
		ax.scatter(self.phi, self.psi, c=colors, picker=True)
		ax.set_xlabel("phi")
		ax.set_xlim(-180, 180)
		ax.set_xticks(range(-180, 181, 60))
		ax.set_ylabel("psi")
		ax.set_ylim(-180, 180)
		ax.set_yticks(range(-180, 181, 60))
		ax.set_title("Ramachandran Plot for %s" % self.mol.name)
		ax.grid(True)
		self.draw()

	def onPick(self, event):
		residues = [ self.residues[i] for i in event.ind ]
		from chimera import selection
		sel = selection.ItemizedSelection(residues)
		selection.setCurrent(sel)

	def _selectionChangedCB(self, trigger, userData, ignore):
		from chimera import selection
		residues = selection.currentResidues()
		self.selectedIndices = []
		for r in residues:
			try:
				n = self.residues.index(r)
			except ValueError:
				pass
			else:
				self.selectedIndices.append(n)
		self._displayData()

	def destroy(self):
		if self.selHandler:
			from chimera import triggers
			triggers.deleteHandler("selection changed",
							self.selHandler)
			self.selHandler = None
		PlotDialog.destroy()

from chimera import openModels, Molecule
m = openModels.list(modelTypes=[Molecule])[0]
RamachandranPlot(m)