from pylab import figure, show, nx
from matplotlib.widgets import Cursor

class DataCursor(Cursor):
    def __init__(self, t, y, ax, useblit=True, **lineprops):
        Cursor.__init__(self, ax, useblit=True, **lineprops)
        self.y = y
        self.t = t
        self.xstr = ''
        self.ystr = ''

    def onmove(self, event):
        """
        we override event.xdata to force it to snap-to nearest data
        item here we assume t is sorted and I'll use searchsorted
        since it is a little faster, but you can plug in your nearest
        neighbor routine, eg to grab the closest x,y point to the
        cursor
        """
        xdata = event.xdata
        ind = nx.searchsorted(self.t, xdata)
        ind = min(len(self.t)-1, ind)
        event.xdata = self.t[ind]
        event.ydata = self.y[ind]
        self.xstr = '%1.3f'%event.xdata
        self.ystr = '%1.3f'%event.ydata
        Cursor.onmove(self, event)
        
    def fmtx(self, x):
        return self.xstr


    def fmty(self, y):
        return self.ystr

fig = figure()
ax1 = fig.add_subplot(211)
ax2 = fig.add_subplot(212, sharex=ax1) # connect x pan/zoom events

t = nx.cumsum(nx.rand(20))
s1 = nx.mlab.rand(len(t))
s2 = nx.mlab.rand(len(t))

ax1.plot(t, s1, 'go')
ax2.plot(t, s2, 'bs')
ax1.set_title("Press 1 for upper cursor and 2 for lower cursor")
cursor1 = DataCursor(t, s1, ax1, useblit=True, color='red', linewidth=2 )

# we'll let the cursor do the toolbarformatting too.
ax1.fmt_xdata = cursor1.fmtx
ax1.fmt_ydata = cursor1.fmty

cursor2 = DataCursor(t, s2, ax2, useblit=True, color='red', linewidth=2 )
ax2.fmt_xdata = cursor2.fmtx
ax2.fmt_ydata = cursor2.fmty


# now we'll control the visibility of the cursor; turn off cursor2 by default
cursor2.visible = False

def keyevent(event):
    cursor1.visible = event.key=='1'
    cursor2.visible = event.key=='2'
    fig.canvas.draw()
    
fig.canvas.mpl_connect('key_press_event', keyevent)
show()