from matplotlib.widgets import Lasso
import matplotlib.mlab
from matplotlib.nxutils import points_inside_poly
from matplotlib.colors import colorConverter
from matplotlib.collections import RegularPolyCollection

from matplotlib.pyplot import figure, show
from numpy import nonzero
from numpy.random import rand

class Datum:
    colorin = colorConverter.to_rgba('red')
    colorout = colorConverter.to_rgba('green')
    def __init__(self, x, y, include=False):
        self.x = x
        self.y = y
        if include: self.color = self.colorin
        else: self.color = self.colorout


class LassoManager:
    def __init__(self, ax, data):
        self.axes = ax
        self.canvas = ax.figure.canvas
        self.data = data
        #the lasso lock boolean is used to tell whether another
        #widget event has priority
        self.lassoLock = False

        self.Nxy = len(data)

        facecolors = [d.color for d in data]
        self.xys = [(d.x, d.y) for d in data]
        fig = ax.figure
        self.collection = RegularPolyCollection(
            fig.dpi, 6, sizes=(100,),
            facecolors=facecolors,
            offsets = self.xys,
            transOffset = ax.transData)

        ax.add_collection(self.collection)

        self.cid = self.canvas.mpl_connect('button_press_event', self.onpress)
        self.cidRelease = self.canvas.mpl_connect('button_release_event', 
self.onrelease)

        self.ind = None

    def callback(self, verts):
        facecolors = self.collection.get_facecolors()
        ind = nonzero(points_inside_poly(self.xys, verts))[0]
        for i in range(self.Nxy):
            if i in ind:
                facecolors[i] = Datum.colorin
            else:
                facecolors[i] = Datum.colorout

        self.canvas.draw_idle()
        self.canvas.widgetlock.release(self.lasso)
        #del self.lasso
        self.ind = ind

    def onpress(self, event):
        if self.canvas.widgetlock.locked(): return
        if event.inaxes is None: return
        self.lasso = Lasso(event.inaxes, (event.xdata, event.ydata), 
self.callback)
        # acquire a lock on the widget drawing
        self.canvas.widgetlock(self.lasso)
        # establish boolean that can be used to release the widgetlock
        self.lassoLock = True

    def onrelease(self, event):
        'on release we reset the press data'
        # test whether the widgetlock was initiated by the lasso
        if self.lassoLock:
            self.canvas.widgetlock.release(self.lasso)
            self.lassoLock = False
        print self.ind


if __name__ == '__main__':

    data = [Datum(*xy) for xy in rand(10, 2)]

    fig = figure()
    ax = fig.add_subplot(111, xlim=(0,10), ylim=(0,10), autoscale_on=True)
    lman = LassoManager(ax, data)

    show()