'''Draw polygon regions of interest (ROIs) in matplotlib images,
similar to Matlab's roipoly function.

See the file example.py for an application. 

Created by Joerg Doepfert 2014 based on code posted by Daniel
Kornhauser.

'''


import numpy as np
import sys
import matplotlib.pyplot as plt
import matplotlib.path as mplPath

class roipoly( object ):

	def __init__(self, fig=[], ax=[], roicolor='r'):
        	if fig == []:  fig = plt.gcf()

		if ax == []: ax = plt.gca()

		self.previous_point = []
        	self.allxpoints = []
        	self.allypoints = []
        	self.start_point = []
        	self.end_point = []
        	self.line = None
        	self.roicolor = roicolor
        	self.fig = fig
        	self.ax = ax
        	#self.fig.canvas.draw()

        	self.__ID1 = self.fig.canvas.mpl_connect(
            	'motion_notify_event', self.__motion_notify_callback)
        	self.__ID2 = self.fig.canvas.mpl_connect(
            	'button_press_event', self.__button_press_callback)

       		if sys.flags.interactive: plt.show(block=False)
        	else: plt.show()

	def getMask(self, currentImage):
        	ny, nx = np.shape(currentImage)
        	poly_verts = [(self.allxpoints[0], self.allypoints[0])]
        	for i in range(len(self.allxpoints)-1, -1, -1):
            		poly_verts.append((self.allxpoints[i], self.allypoints[i]))

        		# Create vertex coordinates for each grid cell...
        		# (<0,0> is at the top left of the grid in this system)
        		x, y = np.meshgrid(np.arange(nx), np.arange(ny))
        		x, y = x.flatten(), y.flatten()
        		points = np.vstack((x,y)).T

        		ROIpath = mplPath.Path(poly_verts)
        		grid = ROIpath.contains_points(points).reshape((ny,nx))
        	return grid
      

    	def __motion_notify_callback(self, event):
        	if event.inaxes:
            		ax = event.inaxes
            		x, y = event.xdata, event.ydata
            	if (event.button == None or event.button == 1) and self.line != None: # Move line around
                	self.line.set_data([self.previous_point[0], x],
                                   [self.previous_point[1], y])
                	self.fig.canvas.draw()


    	def __button_press_callback(self, event):
        	if event.inaxes:
            		x, y = event.xdata, event.ydata
            		ax = event.inaxes
            	if event.button == 1 and event.dblclick == False:  # If you press the left button, single click
                	if self.line == None: # if there is no line, create a line
                    		self.line = plt.Line2D([x, x],
                                           [y, y],
                                           marker='o',
                                           color=self.roicolor)
                    		self.start_point = [x,y]
                    		self.previous_point =  self.start_point
                    		self.allxpoints=[x]
                    		self.allypoints=[y]
                                                
                    		ax.add_line(self.line)
                    		self.fig.canvas.draw()
                    	# add a segment
                	else: # if there is a line, create a segment
                    		self.line = plt.Line2D([self.previous_point[0], x],
                                           [self.previous_point[1], y],
                                           marker = 'o',color=self.roicolor)
                    		self.previous_point = [x,y]
                    		self.allxpoints.append(x)
                    		self.allypoints.append(y)
                                                                                
                    		event.inaxes.add_line(self.line)
                    		self.fig.canvas.draw()
            	elif ((event.button == 1 and event.dblclick==True) or
                  (event.button == 3 and event.dblclick==False)) and self.line != None: # close the loop and disconnect
                	self.fig.canvas.mpl_disconnect(self.__ID1) #joerg
                	self.fig.canvas.mpl_disconnect(self.__ID2) #joerg
                        
                	self.line.set_data([self.previous_point[0],
                                    self.start_point[0]],
                                   [self.previous_point[1],
                                    self.start_point[1]])
                	ax.add_line(self.line)
                	self.fig.canvas.draw()
                	self.line = None
                        
                	if sys.flags.interactive: pass
                	else:
                    		#figure has to be closed so that code can continue
                    		plt.close(self.fig) 


if __name__ == '__main__':

	import h5py


	f = h5py.File('/Users/josh/Documents/Data/Tumor_Sep15/20151005_MRIs.hdf5')
	img = f['meas_MID59_Gre_2D_Gauss_10X1440d_200FOV_2500Hz_FID3330']
	
	plt.imshow( img[0], cmap=plt.cm.binary_r )


	ROI1 = roipoly() #let user draw first ROI

	# show the image with the first ROI
	#plt.imshow(img[0], interpolation='nearest', cmap="Greys")

	msk = ROI1.getMask( img[0] ) 

	plt.imshow( img[0],cmap=plt.cm.binary_r  )	
	plt.imshow( np.ma.masked_equal( (msk*img[0]), 0.0), alpha=0.3 )

	plt.show()
	#np.ma.make_mask( msk )

	

