laydi/fluents/plots.py

import os,sys
from itertools import izip

import pygtk
import gobject
import gtk

import matplotlib
from matplotlib import cm,cbook
from matplotlib.backends.backend_gtkagg import FigureCanvasGTKAgg
from matplotlib.nxutils import points_inside_poly
from matplotlib.axes import Subplot
from matplotlib.figure import Figure
from matplotlib.collections import LineCollection
from matplotlib.patches import Polygon,Rectangle, Circle
from matplotlib.lines import Line2D
from matplotlib.mlab import prctile
import networkx
import scipy

import fluents
import logger

# global active mode. Used by toolbars to communicate correct mode
active_mode = 'default'


class ObjectTable:
    """A 2D table of elements.
    An ObjectTable is a resizable two-dimensional array of objects.
    When resized, it will keep all the elements that fit in the new shape,
    and forget about other elements. When new elements are needed to fill
    out the new shape, they will be created.
    """

    def __init__(self, xsize=0, ysize=0, creator=None):
        """Creates a two-dimensional table of objects.
        @param xsize: Width of the object table.
        @param ysize: Height of the object table.
        @param creator: A function that will be used to instantiate new items.
        """
        self._elements = []
        self._creator = creator or (lambda : None)
        self.xsize = xsize
        self.ysize = ysize
        self.resize(xsize, ysize)

    def resize(self, xsize, ysize):
        """Resizes the table by removing and creating elements as required.
        Makes the ObjectTable the specified size by removing elements that
        will not fit in the new shape and appending elements to fill the
        table. New elements will be created with the creator function passed
        as an argument to __init__.

        @param xsize: The new width.
        @type xsize: int
        @param ysize: The new height.
        @type ysize: int
        """

        # Delete or append new rows
        del self._elements[ysize:]
        new_rows = ysize - len(self._elements)
        self._elements.extend([list() for i in range(new_rows)])

        # Delete or append new columns
        for row in self._elements:
            del row[xsize:]
            new_elems = xsize - len(row)
            row.extend([self._creator() for i in range(new_elems)])

        self.xsize = xsize
        self.ysize = ysize

    def __getitem__(self, index):
        x, y = index
        return self._elements[y][x]

    def __setitem__(self, index, value):
        x, y = index
        self._elements[y][x] = value


class ViewFrame (gtk.Frame):
    """
    A ViewFrame is a gtk container widget that contains a View.
    The ViewFrame is either active or inactive, and belongs to a group of 
    VeiwFrames of which only one can be active at any time.
    """

    def __init__(self, view_frames):
        """Initializes a new ViewFrame.
        @param view_frames: A list of view frames to which the new view
            frame shouls belong. Only one ViewFrame in this list can be
            active at any time.
        """
        gtk.Frame.__init__(self)
        self.focused = False
        self.view_frames = view_frames
        self.empty_view = EmptyView()
        self._button_event = None

        ## Set up a VBox with a label wrapped in an event box.
        label = gtk.Label()
        ebox = gtk.EventBox()
        ebox.add(label)
        vbox = gtk.VBox()
        vbox.pack_start(ebox, expand=False)
        vbox.pack_start(gtk.HSeparator(), expand=False)

        self._ebox_button_event = ebox.connect("button-press-event", 
                                               self._on_button_press_event)
        ## Keep the references for later use.
        self._vbox = vbox
        self._ebox = ebox
        self._view_title = label
        self.add(vbox)
                
        view_frames.append(self)
        if len(view_frames) == 1:
            self.focus()
        else:
            self.focused = True
            self.unfocus()

        # Get dropped views
        self.drag_dest_set(gtk.DEST_DEFAULT_ALL,
                           [("GTK_TREE_MODEL_ROW", gtk.TARGET_SAME_APP, 7)],
                           gtk.gdk.ACTION_LINK)
        self.connect("drag-data-received", self.on_drag_data_received)
       
        # Set view
        self._view = self.empty_view
        self._view.connect("button-press-event", self._on_button_press_event)
        self._vbox.add(self._view)
        self._view_title.set_text(self._view.title)
        self.show_all()
        self._view.show()

    def focus(self):
        """Gets focus and ensures that no other ViewFrame in the same set
        is in focus.
        """
        if self.focused:
            self.emit('focus-changed', self, True)
            return self

        for frame in self.view_frames:
            frame.unfocus()

        self.set_shadow_type(gtk.SHADOW_IN)
        self._ebox.set_state(gtk.STATE_ACTIVE)
        self.focused = True
        self.emit('focus-changed', self, True)
        return self

    def unfocus(self):
        """Removes focus from the ViewFrame. Does nothing if unfocused."""
        if not self.focused:
            return

        self.set_shadow_type(gtk.SHADOW_OUT)
        self._ebox.set_state(gtk.STATE_NORMAL)
        self.focused = False
        self.emit('focus-changed', self, False)

    def set_view(self, view):
        """Set view to view or to empty view if parameter is None
        @param view: An instance of View
        """

        # if None is passed, use empty view
        if view == None:
            view = self.empty_view

        # do nothing if the view is already there
        if view == self._view:
            return
        
        # detach view from current parent
        if view._view_frame:
            view._view_frame.set_view(None)

        # switch which widget we are listening to
        if self._button_event:
            self._view.disconnect(self._button_event)

        self._button_event = view.connect("button-press-event", 
                                          self._on_button_press_event)

        # remove old view, set new view
        if self._view:
            self._view.hide()
            self._vbox.remove(self._view)
            self._view._view_frame = None

        self._view_title.set_text(view.title)
        self._vbox.add(view)
        view.show()
        
        view._view_frame = self
        self._view = view

    def get_view(self):
        """Returns current view, or None if the empty view is set.
        @returns: None if the ViewFrame is currently displaying it's
            EmptyView. Otherwise it returns the currently displeyd View.
        """
        if self._view == self.empty_view:
            return None
        return self._view

    def _on_button_press_event(self, widget, event):
        if not self.focused:
            self.focus()
            
    def on_drag_data_received(self, widget, drag_context, x, y, 
                              selection, info, timestamp):
        treestore, path = selection.tree_get_row_drag_data()
        iter = treestore.get_iter(path)
        obj  = treestore.get_value(iter, 2)

        if isinstance(obj, Plot):
            self.set_view(obj)
            self.focus()
            
        elif isinstance(obj, fluents.dataset.Dataset):
            view = self.get_view()
            if view.is_mappable_with(obj):
                view._update_color_from_dataset(obj)

        elif isinstance(obj, fluents.dataset.Selection):
            view = self.get_view()
            if view.is_mappable_with(obj):
                view.selection_changed(self.current_dim, obj)
            
                

class MainView (gtk.Notebook):
    """The MainView class displays the Views in Fluents.

    MainView, of which there is normally only one instance, contains a 
    gtk.Table that holds all the visible Views, and a single ViewFrame
    that is used for maximizing plots.
    """
    def __init__(self):
        gtk.Notebook.__init__(self)
        self.set_show_tabs(False)
        self.set_show_border(False)

        self._view_frames = []
        self._views = ObjectTable(2, 2, lambda : ViewFrame(self._view_frames))
        self._large_view = ViewFrame(list())
        self.update_small_views()

        for vf in self._view_frames:
            vf.connect('focus-changed', self._on_view_focus_changed)
        
        self.append_page(self._small_views)
        self.append_page(self._large_view)
        self.show()
        self.set_current_page(0)

    def __getitem__(self, x, y):
        return self._views[x, y]

    def update_small_views(self):
        """Creates a new gtk.Table to show the views. Called after changes to
        the _views property"""

        self._small_views = gtk.Table(self._views.ysize, self._views.xsize, True)
        self._small_views.set_col_spacings(4)
        self._small_views.set_row_spacings(4)
        for x in range(self._views.xsize):
            for y in range(self._views.ysize):
                child = self._views[x,y]
                self._small_views.attach(child, x, x+1, y, y+1)
        self._small_views.show_all()

    def get_active_small_view(self):
        """Returns the active ViewFrame in the small views table.
        If a view is maximized, the corresponding ViewFrame in the table
        will be returned, not the active maximized ViewFrame"""
        for vf in self._view_frames:
            if vf.focused:
                return vf
        return None

    def get_active_view_frame(self):
        """Returns the active view frame."""
        if self.get_current_page() == 0:
            return self.get_active_small_view()
        else:
            return self._large_view
            
    def goto_large(self):
        """Maximize the view in the current ViewFrame.
        
        Maximizes the View in the current ViewFrame. The ViewFrame itself is
        not resized or modified, except it's view will be set to None.
        This method will do nothing if a view is currently maximized.
        """
        if self.get_current_page() == 1:
            return

        vf = self.get_active_small_view()
        view = vf.get_view()
        vf.set_view(None)
        self._large_view.set_view(view)
        self.set_current_page(1)

    def goto_small(self):
        """Goes to the small views page.

        The maximized View will be given to the active ViewFrame in the view
        table.
        """
        if self.get_current_page() == 0:
            return

        vf = self.get_active_small_view()
        view = self._large_view.get_view()
        self._large_view.set_view(None)
        vf.set_view(view)
        self.set_current_page(0)

    def insert_view(self, view):
        """Set a view in the currently active ViewFrame.
        @param view: An instance of View.
        """
        if self.get_current_page() == 0:
            vf = self.get_active_small_view()
        else:
            vf = self._large_view
        vf.set_view(view)
       
    def set_all_plots(self, views):
        """Displays all the views in the list, and hides all other views.

        Loops through all ViewFrames from top left to bottom right, and sets
        their currently active View to the next view in the list. After the
        last view is placed in a ViewFrame, the remaining ViewFrames are
        emptied.

        @param views: A list of views to set.
        """
        for y in range(self._views.ysize):
            for x in range(self._views.xsize):
                if views:
                    self._views[x, y].set_view(views.pop())
                else:
                    self._views[x, y].set_view(None)

    def get_active_view_index(self):
        current = self.get_active_small_view()
        for i in range(self._views.xsize):
            for j in range(self._views.ysize):
                if self._views[i, j] == current:
                    return (i, j)
        return None

    def set_active_view_by_index(self, x, y):
        self._views[x, y].focus()

    def move_focus_left(self):
        x, y = self.get_active_view_index()
        if x > 0:
            self.set_active_view_by_index(x-1, y)

    def move_focus_right(self):
        x, y = self.get_active_view_index()
        if x < self._views.xsize-1:
            self.set_active_view_by_index(x+1, y)

    def move_focus_up(self):
        x, y = self.get_active_view_index()
        if y > 0:
            self.set_active_view_by_index(x, y-1)

    def move_focus_down(self):
        x, y = self.get_active_view_index()
        if y < self._views.ysize-1:
            self.set_active_view_by_index(x, y+1)

    def _on_view_focus_changed(self, widget, vf, focused):
        if focused:
            self.emit('view-changed', vf)
    
    def resize_table(self, width, height):
        """Resizes the list of small views.
        """ 
        ## Hide all plots that will be removed from the screen. 
        for x in range(self._views.xsize):
            for y in range(self._views.ysize):
                if x >= width or y >= height:
                    self._views[x, y].set_view(None)
                    self._view_frames.remove(self._views[x, y])

        for x in range(self._views.xsize):
            for y in range(self._views.ysize):
                self._small_views.remove(self._views[x, y])

        self._views.resize(width, height)
        self.update_small_views()
        self.insert_page(self._small_views, gtk.Label("small"), 0)
        self.remove_page(1)
        #self.set_current_page(0)
        self.goto_small()


class View (gtk.Frame):
    """The base class of everything that is shown in the center view of
    fluents.

    Most views should rather subclass Plot, which automatically handles 
    freezing, creates a toolbar, and sets up matplotlib Figure and Canvas 
    objects.
    """

    def __init__(self, title):
        gtk.Frame.__init__(self)
        self.title = title
        self.set_shadow_type(gtk.SHADOW_NONE)
        self._view_frame = None

    def get_toolbar(self):
        return None

    def is_mappable_with(self, dataset):
        """Override in individual plots."""
        return False


class EmptyView (View):
    """EmptyView is shown in ViewFrames that are unused."""
    def __init__(self):
        View.__init__(self, 'Empty view')
        self.set_label(None)
        label = gtk.Label('No view')
        ebox = gtk.EventBox()
        ebox.add(label)
        self.add(ebox)
        label.show()
        ebox.show()
        self.show()

    
class Plot (View):
    def __init__(self, title):
        View.__init__(self, title)
        logger.log('debug', 'plot %s init' %title)
        self.selection_listener = None
        self.current_dim = None
        self._current_selection = None
        self._frozen = False
        self._init_mpl()
        
    def _init_mpl(self):    
        # init matplotlib related stuff 
        self._background = None
        self._use_blit = False
        self.fig = Figure()
        self.canvas = FigureCanvasGTKAgg(self.fig)
        self.axes = self.fig.gca()
        self._toolbar = PlotToolbar(self)
        self.canvas.add_events(gtk.gdk.ENTER_NOTIFY_MASK)
        self.add(self.canvas)
        self.canvas.show()
    
    def set_frozen(self, frozen):
        """A frozen plot will not be updated when the current
        selection is changed."""
        self._frozen = frozen
        if not frozen and self._current_selection != None:
            self.set_current_selection(self._current_selection)
    
    def get_title(self):
        return self.title

    def get_toolbar(self):
        return self._toolbar
    
    def selection_changed(self, dim_name, selection):
        """ Selection observer handle.
        
        A selection change in a plot is only drawn if:
        1.) plot is sensitive to selections (not freezed)
        2.) plot is visible (has a view)
        3.) the selections dim_name is the plot's dimension.
        """
        
        self._current_selection = selection
        if self._frozen \
               or not self.get_property('visible') \
               or self.current_dim != dim_name:
            return

        self.set_current_selection(selection)

    def set_selection_listener(self, listener):
        """Allow project to listen to selections.

        The selection will propagate back to all plots through the
        selection_changed() method. The listener will be called as
        listener(dimension_name, ids).
        """
        self.selection_listener = listener

    def update_selection(self, ids, key=None):
        """Returns updated current selection from ids.
        If a key is pressed we use the appropriate mode.

        key map:
        shift : union
        control : intersection

        """
        if key == 'shift':
            ids = set(ids).union(self._current_selection[self.current_dim])
        elif key == 'control':
            ids = set(ids).intersection(self._current_selection[self.current_dim])
        return ids

    def set_current_selection(self, selection):
        """Called whenever the plot should change the selection.

        This method is a dummy method, so that specialized plots that have 
        no implemented selection can ignore selections alltogether.
        """
        pass

    def rectangle_select_callback(self, *args):
        """Overrriden in subclass."""
        if hasattr(self, 'canvas'):
            self.canvas.draw()

    def lasso_select_callback(self, *args):
        """Overrriden in subclass."""
        if hasattr(self, 'canvas'):
            self.canvas.draw()
    
    def get_index_from_selection(self, dataset, selection):
        """Returns the index vector of current selection in given dim."""
        if not selection: return []
        ids = selection.get(self.current_dim, []) # current identifiers
        if not ids : return []
        return dataset.get_indices(self.current_dim, ids)


class LineViewPlot(Plot):
    """Line view plot with percentiles.

    A line view of vectors across a specified dimension of input dataset. 
    No selection interaction is defined.
    Only support for 2d-arrays.
    input:
           -- major_axis : dim_number for line dim (see scipy.ndarray for axis def.)
           -- minor_axis : needs definition only for higher order arrays

    fixme: slow
    """
    def __init__(self, dataset, major_axis=1, minor_axis=None, name="Line view"):
        Plot.__init__(self, name)
        self.dataset = dataset
        self._data = dataset.asarray()
        if len(self._data.shape)==2 and not minor_axis:
            minor_axis = major_axis - 1
        self.major_axis = major_axis
        self.minor_axis = minor_axis
        self.current_dim = self.dataset.get_dim_name(major_axis)
        
        #initial draw
        self.line_coll = None
        self.line_segs = []
        x_axis = scipy.arange(self._data.shape[minor_axis])
        for xi in range(self._data.shape[major_axis]):
            yi = self._data.take([xi], major_axis).ravel()
            self.line_segs.append([(xx,yy) for xx,yy in izip(x_axis, yi)])
        
        # draw background
        self._set_background(self.axes)
        
        # Disable selection modes
        self._toolbar.freeze_button.set_sensitive(False)
        self._toolbar.set_mode_sensitive('select', False)
        self._toolbar.set_mode_sensitive('lassoselect', False)
        
    def _set_background(self, ax):
        """Add three patches representing [min max],[5,95] and [25,75] percentiles, and a line at the median.
        """
        if self._data.shape[self.minor_axis]<10:
            return 
        # settings
        patch_color = 'b' #blue
        patch_lw = 0 #no edges
        patch_alpha = .15 # transparancy
        median_color = 'b' #blue
        median_width = 1.5 #linewidth
        percentiles = [0, 5, 25, 50, 75, 100]

        # ordinate
        xax = scipy.arange(self._data.shape[self.minor_axis])

        #vertices
        verts_0 = [] #100,0
        verts_1 = [] # 90,10
        verts_2 = [] # 75,25
        med = []
        # add top vertices the low vertices (do i need an order?)#background
        for i in xax:
            prct = prctile(self._data.take([i], self.minor_axis), percentiles)
            verts_0.append((i, prct[0]))
            verts_1.append((i, prct[1]))
            verts_2.append((i, prct[2]))
            med.append(prct[3])
        for i in xax[::-1]:
            prct = prctile(self._data.take([i], self.minor_axis), percentiles)
            verts_0.append((i, prct[-1]))
            verts_1.append((i, prct[-2]))
            verts_2.append((i, prct[-3]))

        # make polygons from vertices
        bck0 = Polygon(verts_0, alpha=patch_alpha, lw=patch_lw,
                       facecolor=patch_color)
        bck1 = Polygon(verts_1, alpha=patch_alpha, lw=patch_lw,
                       facecolor=patch_color)
        bck2 = Polygon(verts_2, alpha=patch_alpha, lw=patch_lw,
                       facecolor=patch_color)

        # add polygons to axes
        ax.add_patch(bck0)
        ax.add_patch(bck1)
        ax.add_patch(bck2)
        # median line
        ax.plot(xax, med, median_color, linewidth=median_width)

        
    def set_current_selection(self, selection):
        """Draws the current selection.
        """
        index = self.get_index_from_selection(self.dataset, selection)

        if self.line_coll:
            self.axes.collections.remove(self.line_coll)
        segs = [self.line_segs[i] for i in index]
        self.line_coll = LineCollection(segs, colors=(1,0,0,1))
        self.axes.add_collection(self.line_coll)
        
        #draw
        if self._use_blit:
            if self._background is None:
                self._background = self.canvas.copy_from_bbox(self.axes.bbox)
                self.canvas.restore_region(self._background)
            self.axes.draw_artist(self.lines)
            self.canvas.blit()
        else:
            self.canvas.draw()
            

class ScatterMarkerPlot(Plot):
    """The ScatterMarkerPlot is faster than regular scatterplot, but
       has no color and size options."""
    
    def __init__(self, dataset_1, dataset_2, id_dim, sel_dim,
                 id_1, id_2, s=6, name="Scatter plot"):
        Plot.__init__(self, name)
        self.current_dim = id_dim
        self.dataset_1 = dataset_1
        self.ms = s
        x_index = dataset_1[sel_dim][id_1]
        y_index = dataset_2[sel_dim][id_2]
        self.xaxis_data = dataset_1._array[:, x_index]
        self.yaxis_data = dataset_2._array[:, y_index]

        # init draw
        self._selection_line = None
        self.line = self.axes.plot(self.xaxis_data, self.yaxis_data, 'o', markeredgewidth=0, markersize=s)
        self.axes.axhline(0, color='k', lw=1., zorder=1)
        self.axes.axvline(0, color='k', lw=1., zorder=1)

    def rectangle_select_callback(self, x1, y1, x2, y2, key):
        ydata = self.yaxis_data
        xdata = self.xaxis_data

        # find indices of selected area
        if x1>x2:
            x1, x2 = x2, x1
        if y1>y2:
            y1, y2 = y2, y1
        assert x1<=x2
        assert y1<=y2
        index = scipy.nonzero((xdata>x1) & (xdata<x2) & (ydata>y1) & (ydata<y2))[0]
        ids = self.dataset_1.get_identifiers(self.current_dim, index)
        ids = self.update_selection(ids, key)
        self.selection_listener(self.current_dim, ids)

    def lasso_select_callback(self, verts, key=None):
        xys = scipy.c_[self.xaxis_data[:,scipy.newaxis], self.yaxis_data[:,scipy.newaxis]]
        index = scipy.nonzero(points_inside_poly(xys, verts))[0]
        ids = self.dataset_1.get_identifiers(self.current_dim, index)
        ids = self.update_selection(ids, key)
        self.selection_listener(self.current_dim, ids)
    
    def set_current_selection(self, selection):
        #remove old selection
        if self._selection_line:
           self.axes.lines.remove(self._selection_line)
        index = self.get_index_from_selection(self.dataset_1, selection)
        if len(index)==0:
            # no selection
            self.canvas.draw()
            self._selection_line = None
            return
        
        xdata_new = self.xaxis_data.take(index) #take data
        ydata_new = self.yaxis_data.take(index)
        self._selection_line = Line2D(xdata_new, ydata_new
                                      ,marker='o', markersize=self.ms,
                                      linewidth=0, markerfacecolor='r',
                                      markeredgewidth=1.0)
        self.axes.add_line(self._selection_line)

        if self._use_blit:
            if self._background is None:
                self._background = self.canvas.copy_from_bbox(self.axes.bbox)
            self.canvas.restore_region(self._background)
            if self.selection_line:
                self.axes.draw_artist(self._selection_line)
            self.canvas.blit()
        else:
            self.canvas.draw()


class ScatterPlot(Plot):
    """The ScatterPlot is slower than scattermarker, but has size option."""
    def __init__(self, dataset_1, dataset_2, id_dim, sel_dim, id_1, id_2, c='b', s=30, sel_dim_2=None, name="Scatter plot"):

        Plot.__init__(self, name)
        self.dataset_1 = dataset_1
        self.s = s
        self.c = c
        self.current_dim = id_dim
        
        x_index = dataset_1[sel_dim][id_1]
        if sel_dim_2:
            y_index = dataset_2[sel_dim_2][id_2]
        else:
            y_index = dataset_2[sel_dim][id_2]
        self.xaxis_data = dataset_1._array[:, x_index]
        self.yaxis_data = dataset_2._array[:, y_index]

        # init draw
        self.init_draw()

        # signals to enable correct use of blit
        self.connect('zoom-changed', self.onzoom)
        self.connect('pan-changed', self.onpan)
        self.need_redraw = False
        self.canvas.mpl_connect('resize_event', self.onresize)
        
    def onzoom(self, widget, mode):
        logger.log('notice', 'Zoom in widget: %s' %widget)
        self.clean_redraw()
        
    def onpan(self, widget, mode):
        logger.log('notice', 'Pan in widget: %s' %widget)
        self.clean_redraw()
        
    def onresize(self, widget):
        logger.log('notice', 'resize event')
        self.clean_redraw()

    def clean_redraw(self):
        if self._use_blit == True:
            logger.log('notice', 'blit -> clean redraw ')
            self.set_current_selection(None)
            self._background = self.canvas.copy_from_bbox(self.axes.bbox)
            self.set_current_selection(self._current_selection)
        else:
            self._background = None
        
    def init_draw(self):
        lw = scipy.zeros(self.xaxis_data.shape)
        self.sc = self.axes.scatter(self.xaxis_data, self.yaxis_data,
                                  s=self.s, c=self.c, linewidth=lw, zorder=3)
        self.axes.axhline(0, color='k', lw=1., zorder=1)
        self.axes.axvline(0, color='k', lw=1., zorder=1)
        self.selection_collection = self.axes.scatter(self.xaxis_data,
                                                      self.yaxis_data,
                                                      alpha=0,
                                                      c='w',s=self.s,
                                                      linewidth=0,
                                                      zorder=4)
        self._background = self.canvas.copy_from_bbox(self.axes.bbox)
    
    def is_mappable_with(self, obj):
        """Returns True if dataset/selection is mappable with this plot.
        """
        if isinstance(obj, fluents.dataset.Dataset):
            if self.current_dim in obj.get_dim_name() \
                   and obj.asarray().shape[0] == self.xaxis_data.shape[0]:
                return True
        
        elif isinstance(obj, fluents.dataset.Selection):
            if self.current_dim in obj.get_dim_name():
                print "Selection is mappable"
                return True

        else:
            return False
        
    def _update_color_from_dataset(self, data):
        """Updates the facecolors from a dataset.
        """
        array = data.asarray()
        #only support for 2d-arrays:
        try:
            m, n = array.shape
        except:
            raise ValueError, "No support for more than 2 dimensions."
        # is dataset a vector or matrix?
        if not n==1:
            # we have a category dataset
            if isinstance(data, fluents.dataset.CategoryDataset):
                map_vec = scipy.dot(array, scipy.diag(scipy.arange(n))).sum(1)
            else:
                map_vec = array.sum(1)
        else:
            map_vec = array.ravel()

        # update facecolors
        self.sc.set_array(map_vec)
        self.sc.set_clim(map_vec.min(), map_vec.max())
        self.sc.update_scalarmappable() #sets facecolors from array
        self.canvas.draw()
                    
    def rectangle_select_callback(self, x1, y1, x2, y2, key):
        ydata = self.yaxis_data
        xdata = self.xaxis_data

        # find indices of selected area
        if x1>x2:
            x1, x2 = x2, x1
        if y1>y2:
            y1, y2 = y2, y1
        assert x1<=x2
        assert y1<=y2

        index = scipy.nonzero((xdata>x1) & (xdata<x2) & (ydata>y1) & (ydata<y2))[0]
        ids = self.dataset_1.get_identifiers(self.current_dim, index)
        ids = self.update_selection(ids, key)
        self.selection_listener(self.current_dim, ids)

    def lasso_select_callback(self, verts, key=None):
        xys = scipy.c_[self.xaxis_data[:,scipy.newaxis], self.yaxis_data[:,scipy.newaxis]]
        index = scipy.nonzero(points_inside_poly(xys, verts))[0]
        ids = self.dataset_1.get_identifiers(self.current_dim, index)
        ids = self.update_selection(ids, key)
        self.selection_listener(self.current_dim, ids)
        
    def set_current_selection(self, selection):
        linewidth = scipy.zeros(self.xaxis_data.shape, 'f')
        index = self.get_index_from_selection(self.dataset_1, selection)
        if len(index) > 0:
            linewidth.put(1, index)
        self.selection_collection.set_linewidth(linewidth)
            
        if self._use_blit and len(index)>0 :
            if self._background is None:
                    self._background = self.canvas.copy_from_bbox(self.axes.bbox)
            self.canvas.restore_region(self._background)
            self.axes.draw_artist(self.selection_collection)
            self.canvas.blit()
        else:
            self.canvas.draw()
        
        
class ImagePlot(Plot):
    def __init__(self, dataset, **kw):
        Plot.__init__(self, kw.get('name', 'Image Plot'))
        self.dataset = dataset

        # Initial draw
        self.axes.grid(False)
        self.axes.imshow(dataset.asarray(), interpolation='nearest')
        self.axes.axis('tight')

        # Disable selection modes
        self._toolbar.freeze_button.set_sensitive(False)
        self._toolbar.set_mode_sensitive('select', False)
        self._toolbar.set_mode_sensitive('lassoselect', False)

    
class HistogramPlot(Plot):
    """ Histogram plot.
    If dataset is 1-dim the current_dim is set and selections may
    be performed. For dataset> 1.dim the histogram is over all values
    and selections are not defined,"""
    
    def __init__(self, dataset, **kw):
        Plot.__init__(self, kw['name'])
        self.dataset = dataset
        self._data = dataset.asarray()

        # If dataset is 1-dim we may do selections
        if dataset.shape[0]==1:
            self.current_dim = dataset.get_dim_name(1)
        if dataset.shape[1]==1:
            self.current_dim = dataset.get_dim_name(0)
        # Initial draw
        self.axes.grid(False)
        bins = min(self._data.size, 20)
        count, lims, self.patches = self.axes.hist(self._data, bins=bins)

        # Add identifiers to the individual patches
        if self.current_dim != None:
            for i, patch in enumerate(self.patches):
                if i==len(self.patches)-1:
                    end_lim = self._data.max() + 1
                else:
                    end_lim = lims[i+1]
                bool_ind = scipy.bitwise_and(self._data>=lims[i],
                                             self._data<=end_lim)
                patch.index = scipy.where(bool_ind)[0]

        if self.current_dim==None:
            # Disable selection modes
            self._toolbar.freeze_button.set_sensitive(False)
            self._toolbar.set_mode_sensitive('select', False)
            self._toolbar.set_mode_sensitive('lassoselect', False)
    
    def rectangle_select_callback(self, x1, y1, x2, y2, key):
        if self.current_dim == None: return
        # make (x1, y1) the lower left corner
        if x1>x2:
            x1, x2 = x2, x1
        if y1>y2:
            y1, y2 = y2, y1

        self.active_patches = []
        for patch in self.patches:
            xmin = patch.xy[0]
            xmax = xmin + patch.get_width()
            ymin, ymax = 0, patch.get_height() 
            if xmax>x1 and xmin<x2 and (ymax> y2 or ymax>y1):
                self.active_patches.append(patch)
        if not self.active_patches: return

        ids = set()
        for patch in self.active_patches:
            ids.update(self.dataset.get_identifiers(self.current_dim,
                                                    patch.index))
        ids = self.update_selection(ids, key)
        self.selection_listener(self.current_dim, ids)

    def lasso_select_callback(self, verts, key):
        if self.current_dim == None: return
        self.active_patches = []
        for patch in self.patches:
            if scipy.any(points_inside_poly(verts, patch.get_verts())):
                self.active_patches.append(patch)
        if not self.active_patches: return
        ids = set()
        for patch in self.active_patches:
            ids.update(self.dataset.get_identifiers(self.current_dim,
                                                    patch.index))
        ids = self.update_selection(ids, key)
        self.selection_listener(self.current_dim, ids)

    def set_current_selection(self, selection):
        index = self.get_index_from_selection(self.dataset, selection)
        for patch in self.patches:
            patch.set_facecolor('b')
        for patch in self.patches:
            if scipy.intersect1d(patch.index, index).size>1:
                patch.set_facecolor('r')
        self.canvas.draw()


class BarPlot(Plot):
    """Bar plot.

    Ordinary bar plot for (column) vectors.
    For matrices there is one color for each row.
    """
    def __init__(self, dataset, **kw):
        Plot.__init__(self, kw.get('name', 'Bar Plot'))
        self.dataset = dataset
        
        # Initial draw
        self.axes.grid(False)
        n, m = dataset.shape
        if m>1:
            sm = cm.ScalarMappable()
            clrs = sm.to_rgba(range(n))
            for i, row in enumerate(dataset.asarray()):
                left = scipy.arange(i+1, m*n+1, n)
                height = row
                color = clrs[i]
                c = (color[0], color[1], color[2])
                self.axes.bar(left, height,color=c)
        else:
            height = dataset.asarray().ravel()
            left = scipy.arange(1, n, 1)
            self.axes.bar(left, height)

        # Disable selection modes
        self._toolbar.freeze_button.set_sensitive(False)
        self._toolbar.set_mode_sensitive('select', False)
        self._toolbar.set_mode_sensitive('lassoselect', False)
        

class NetworkPlot(Plot):
    def __init__(self, dataset, pos=None, nodecolor='b', nodesize=40,
                 prog='neato', with_labels=False, name='Network Plot'):

        Plot.__init__(self, name)
        self.dataset = dataset
        self.graph = dataset.asnetworkx()
        self._prog = prog
        self._pos = pos
        self._nodesize = nodesize
        self._nodecolor = nodecolor
        self._with_labels = with_labels
        
        self.current_dim = self.dataset.get_dim_name(0)

        if not self._pos:
            self._pos = networkx.graphviz_layout(self.graph, self._prog)
        self._xy = scipy.asarray([self._pos[node] for node in self.dataset.get_identifiers(self.current_dim, sorted=True)])
        self.xaxis_data = self._xy[:,0]
        self.yaxis_data = self._xy[:,1]

        # Initial draw
        self.default_props = {'nodesize': 50, 'nodecolor':'gray'}
        self.node_collection = None
        self.edge_collection = None
        self.node_labels = None
        lw = scipy.zeros(self.xaxis_data.shape)
        self.node_collection = self.axes.scatter(self.xaxis_data, self.yaxis_data,
                                                 s=self._nodesize,
                                                 c=self._nodecolor,
                                                 linewidth=lw,
                                                 zorder=3)
        # selected nodes is a transparent graph that adjust edge-visibility
        # according to the current selection
        self.selected_nodes = self.axes.scatter(self.xaxis_data,
                                                self.yaxis_data,
                                                s=self._nodesize,
                                                c=self._nodecolor,
                                                linewidth=lw,
                                                zorder=4,
                                                alpha=0)
        
        self.edge_collection = networkx.draw_networkx_edges(self.graph,
                                                            self._pos,
                                                            ax=self.axes,
                                                            edge_color='gray')
        if self._with_labels:
            self.node_labels = networkx.draw_networkx_labels(self.graph,
                                                            self._pos,
                                                            ax=self.axes)
        
        # remove axes, frame and grid
        self.axes.set_xticks([])
        self.axes.set_yticks([])
        self.axes.grid(False)
        self.axes.set_frame_on(False)
        self.fig.subplots_adjust(left=0, right=1, bottom=0, top=1)

    def rectangle_select_callback(self, x1, y1, x2, y2, key):
        ydata = self.yaxis_data
        xdata = self.xaxis_data
        
        # find indices of selected area
        if x1>x2:
            x1, x2 = x2, x1
        if y1>y2:
            y1, y2 = y2, y1
        assert x1<=x2
        assert y1<=y2
        
        index = scipy.nonzero((xdata>x1) & (xdata<x2) & (ydata>y1) & (ydata<y2))[0]
        ids = self.dataset.get_identifiers(self.current_dim, index)
        ids = self.update_selection(ids, key)
        self.selection_listener(self.current_dim, ids)

    def lasso_select_callback(self, verts, key=None):
        xys = scipy.c_[self.xaxis_data[:,scipy.newaxis], self.yaxis_data[:,scipy.newaxis]]
        index = scipy.nonzero(points_inside_poly(xys, verts))[0]
        ids = self.dataset.get_identifiers(self.current_dim, index)
        ids = self.update_selection(ids, key)
        self.selection_listener(self.current_dim, ids)
        
    def set_current_selection(self, selection):
        linewidth = scipy.zeros(self.xaxis_data.shape, 'f')
        index = self.get_index_from_selection(self.dataset, selection)
        if len(index) > 0:
            linewidth.put(2, index)
        self.selected_nodes.set_linewidth(linewidth)
        self.canvas.draw()


class VennPlot(Plot):
    def __init__(self, name="Venn diagram"):
        Plot.__init__(self, name)
        
        self._init_bck()
        # init draw
        for c in self._venn_patches:
            self.axes.add_patch(c) 
        for mrk in self._markers:
            self.axes.add_patch(mrk)
        self.axes.set_xlim([-3, 3])
        self.axes.set_ylim([-2.5, 3.5])
        self._last_active = set()
        self.axes.set_xticks([])
        self.axes.set_yticks([])
        self.axes.axis('equal')
        self.axes.grid(False)
        self.axes.set_frame_on(False)
        self.fig.subplots_adjust(left=0, right=1, bottom=0, top=1)

    def _init_bck(self):
        res = 50
        a = .5
        r = 1.5
        mr = .2
        self.c1 = c1 =  Circle((-1,0), radius=r, alpha=a, facecolor='b')
        self.c2 = c2 = Circle((1,0), radius=r, alpha=a, facecolor='r')
        self.c3 = c3 = Circle((0, scipy.sqrt(3)), radius=r, alpha=a, facecolor='g')

        self.c1marker = Circle((-1.25, -.25), radius=mr, facecolor='y', alpha=0)
        self.c2marker = Circle((1.25, -.25), radius=mr, facecolor='y', alpha=0)
        self.c3marker = Circle((0, scipy.sqrt(3)+.25), radius=mr, facecolor='y', alpha=0)
        self.c1c2marker = Circle((0, -.15), radius=mr, facecolor='y', alpha=0)

        self.c1c3marker = Circle((-scipy.sqrt(2)/2, 1), radius=mr, facecolor='y', alpha=0)
        self.c2c3marker = Circle((scipy.sqrt(2)/2, 1), radius=mr, facecolor='y', alpha=0)
        self.c1c2c3marker = Circle((0, .6), radius=mr, facecolor='y', alpha=0)

        c1.elements = set(['a', 'b', 'c', 'f'])
        c2.elements = set(['a', 'c', 'd', 'e']) 
        c3.elements = set(['a', 'e', 'f', 'g'])
        self.active_elements = set()
        self.all_elements = c1.elements.union(c2.elements).union(c3.elements)

        c1.active = False
        c2.active = False
        c3.active = False

        c1.name = 'Blue'
        c2.name = 'Red'
        c3.name = 'Green'
        
        self._venn_patches = [c1, c2, c3]
        self._markers = [self.c1marker, self.c2marker, self.c3marker,
                         self.c1c2marker, self.c1c3marker,
                         self.c2c3marker, self.c1c2c3marker]

        self._tot_label = 'Tot: ' + str(len(self.all_elements))
        self._sel_label = 'Sel: ' + str(len(self.active_elements))
        self._legend = self.axes.legend((self._tot_label, self._sel_label),
                                       loc='upper right')
        
    def set_selection(self, selection, patch=None):
        if patch:
            patch.selection = selection
        else:
            selection_set = False
            for patch in self._venn_patches:
                if len(patch.elements)==0:
                    patch.elements = selection
                    selection_set = True
            if not selection_set:
                self.venn_patches[0].elements = selection
    
    def lasso_select_callback(self, verts, key=None):
        if verts==None:
            verts = (self._event.xdata, self._event.ydata)
        if key!='shift':
            for m in self._markers:
                m.set_alpha(0)
        
        self._patches_within_verts(verts, key)
        active = [i.active for i in self._venn_patches]
        if active==[True, False, False]:
            self.c1marker.set_alpha(1)
            self.active_elements = self.c1.elements.difference(self.c2.elements.union(self.c3.elements))
        elif active== [False, True, False]:
            self.c2marker.set_alpha(1)
            self.active_elements = self.c2.elements.difference(self.c1.elements.union(self.c3.elements))
        elif active== [False, False, True]:
            self.c3marker.set_alpha(1)
            self.active_elements = self.c3.elements.difference(self.c2.elements.union(self.c1.elements))
        elif active==[True, True, False]:
            self.c1c2marker.set_alpha(1)
            self.active_elements = self.c1.elements.intersection(self.c2.elements)
        elif active==[True, False, True]:
            self.c1c3marker.set_alpha(1)
            self.active_elements = self.c1.elements.intersection(self.c3.elements)
        elif active==[False, True, True]:
            self.c2c3marker.set_alpha(1)
            self.active_elements = self.c2.elements.intersection(self.c3.elements)
        elif active==[True, True, True]:
            self.c1c2c3marker.set_alpha(1)
            self.active_elements = self.c1.elements.intersection(self.c3.elements).intersection(self.c2.elements)

        if key=='shift':
            self.active_elements = self.active_elements.union(self._last_active)
        self._last_active = self.active_elements.copy()
        self._sel_label = 'Sel: ' + str(len(self.active_elements))
        self._legend.texts[1].set_text(self._sel_label)
        self.axes.figure.canvas.draw()

    def rectangle_select_callback(self, x1, y1, x2, y2, key):
        verts = [(x1, y1), (x2, y2)]
        if key!='shift':
            for m in self._markers:
                m.set_alpha(0)
        
        self._patches_within_verts(verts, key)
        active = [i.active for i in self._venn_patches]
        if active==[True, False, False]:
            self.c1marker.set_alpha(1)
            self.active_elements = self.c1.elements.difference(self.c2.elements.union(self.c3.elements))
        elif active== [False, True, False]:
            self.c2marker.set_alpha(1)
            self.active_elements = self.c2.elements.difference(self.c1.elements.union(self.c3.elements))
        elif active== [False, False, True]:
            self.c3marker.set_alpha(1)
            self.active_elements = self.c3.elements.difference(self.c2.elements.union(self.c1.elements))
        elif active==[True, True, False]:
            self.c1c2marker.set_alpha(1)
            self.active_elements = self.c1.elements.intersection(self.c2.elements)
        elif active==[True, False, True]:
            self.c1c3marker.set_alpha(1)
            self.active_elements = self.c1.elements.intersection(self.c3.elements)
        elif active==[False, True, True]:
            self.c2c3marker.set_alpha(1)
            self.active_elements = self.c2.elements.intersection(self.c3.elements)
        elif active==[True, True, True]:
            self.c1c2c3marker.set_alpha(1)
            self.active_elements = self.c1.elements.intersection(self.c3.elements).intersection(self.c2.elements)

        if key=='shift':
            self.active_elements = self.active_elements.union(self._last_active)
        self._last_active = self.active_elements.copy()
        self._sel_label = 'Sel: ' + str(len(self.active_elements))
        self._legend.texts[1].set_text(self._sel_label)
        self.axes.figure.canvas.draw()
        
    def _patches_within_verts(self, verts, key):
        xy = scipy.array(verts).mean(0)
        for venn_patch in self._venn_patches:
            venn_patch.active = False
            if self._distance(venn_patch.center,xy)<venn_patch.radius:
                venn_patch.active = True

    def _distance(self, (x1,y1),(x2,y2)):
        return scipy.sqrt( (x2-x1)**2 + (y2-y1)**2 )


class PlotMode:
    """A PlotMode object corresponds to a mouse mode in a plot.

    When a mode is selected in the toolbar, the PlotMode corresponding
    to the toolbar button is activated by calling setup(ax) for the axis
    system ax.
    """
    def __init__(self, plot, name, tooltip, image_file):
        self.name = name
        self.tooltip = tooltip
        self.image_file = image_file
        self.plot = plot
        self.canvas = plot.canvas

    def get_icon(self):
        """Returns the icon for the PlotMode"""
        image = gtk.Image()
        image.set_from_pixbuf(fluents.icon_factory.get(self.image_file))
        return image
        
    def activate(self):
        """Subclasses of PlotMode should do their initialization here.

        The activate method is called when a mode is activated, and is
        used primarily to set up callback functions on events in the
        canvas.
        """
        pass

    def deactivate(self):
        """Subclasses of PlotMode should do their cleanup here.

        The deactivate method is primarily by subclasses of PlotMode to
        remove any callbacks they might have on the matplotlib canvas.
        """
        pass


class DefaultPlotMode(PlotMode):
    def __init__(self, plot):
        PlotMode.__init__(self, plot, 'default', 'Default mode', 'cursor')


class PanPlotMode(PlotMode):
    def __init__(self, plot):
        PlotMode.__init__(self, plot, 'pan',
                          'Pan axes with left mouse, zoom with right',
                          'move')

        # Holds handler IDs for callbacks.
        self._button_press = None
        self._button_release = None
        self._motion_notify = None
        self._xypress = None
        self._button_pressed = None

    def activate(self):
        self._button_press = self.canvas.mpl_connect(
            'button_press_event', self._on_button_press)
        self._button_release = self.canvas.mpl_connect(
            'button_release_event', self._on_button_release)

    def deactivate(self):
        if self._button_press:
            self.canvas.mpl_disconnect(self._button_press)

        if self._button_release:
            self.canvas.mpl_disconnect(self._button_release)

    def _on_button_press(self, event):

        if event.button == 1:
            self._button_pressed = 1
        elif  event.button == 3:
            self._button_pressed = 3
        else:
            self._button_pressed=None
            return

        x, y = event.x, event.y

        # push the current view to define home if stack is empty
        # if self._views.empty(): self.push_current()

        self._xypress=[]
        for i, a in enumerate(self.canvas.figure.get_axes()):
            if x is not None and y is not None and a.in_axes(x, y) \
                   and a.get_navigate():
                xmin, xmax = a.get_xlim()
                ymin, ymax = a.get_ylim()
                lim = xmin, xmax, ymin, ymax
                self._xypress.append((x, y, a, i, lim,a.transData.deepcopy()))
                self.canvas.mpl_disconnect(self._motion_notify)
                
                cid = self.canvas.mpl_connect('motion_notify_event',
                                              self._on_motion_notify)
                self._motion_notify = cid

    def _on_motion_notify(self, event):
        """The drag callback in pan/zoom mode"""

        def format_deltas(event, dx, dy):
            """Returns the correct dx and dy based on the modifier keys"""
            if event.key=='control':
                if(abs(dx)>abs(dy)):
                    dy = dx
                else:
                    dx = dy
            elif event.key=='x':
                dy = 0
            elif event.key=='y':
                dx = 0
            elif event.key=='shift':
                if 2*abs(dx) < abs(dy):
                    dx=0
                elif 2*abs(dy) < abs(dx):
                    dy=0
                elif(abs(dx)>abs(dy)):
                    dy=dy/abs(dy)*abs(dx)
                else:
                    dx=dx/abs(dx)*abs(dy)
            return (dx,dy)

        for cur_xypress in self._xypress:
            lastx, lasty, a, ind, lim, trans = cur_xypress
            xmin, xmax, ymin, ymax = lim

            #safer to use the recorded button at the press than current button:
            #multiple button can get pressed during motion...
            if self._button_pressed==1:
                lastx, lasty = trans.inverse_xy_tup( (lastx, lasty) )
                x, y = trans.inverse_xy_tup( (event.x, event.y) )
                if a.get_xscale()=='log':
                    dx=1-lastx/x
                else:
                    dx=x-lastx
                if a.get_yscale()=='log':
                    dy=1-lasty/y
                else:
                    dy=y-lasty

                dx,dy=format_deltas(event,dx,dy)

                if a.get_xscale()=='log':
                    xmin *= 1-dx
                    xmax *= 1-dx
                else:
                    xmin -= dx
                    xmax -= dx
                if a.get_yscale()=='log':
                    ymin *= 1-dy
                    ymax *= 1-dy
                else:
                    ymin -= dy
                    ymax -= dy
                    
            elif self._button_pressed==3:
                try:
                    dx=(lastx-event.x)/float(a.bbox.width())
                    dy=(lasty-event.y)/float(a.bbox.height())
                    dx,dy=format_deltas(event,dx,dy)
                    if a.get_aspect() != 'auto':
                        dx = 0.5*(dx + dy)
                        dy = dx
                    alphax = pow(10.0,dx)
                    alphay = pow(10.0,dy)
                    lastx, lasty = trans.inverse_xy_tup( (lastx, lasty) )
                    if a.get_xscale()=='log':
                        xmin = lastx*(xmin/lastx)**alphax
                        xmax = lastx*(xmax/lastx)**alphax
                    else:
                        xmin = lastx+alphax*(xmin-lastx)
                        xmax = lastx+alphax*(xmax-lastx)
                    if a.get_yscale()=='log':
                        ymin = lasty*(ymin/lasty)**alphay
                        ymax = lasty*(ymax/lasty)**alphay
                    else:
                        ymin = lasty+alphay*(ymin-lasty)
                        ymax = lasty+alphay*(ymax-lasty)
                        
                except OverflowError:
                    warnings.warn('Overflow while panning')
                    return

            a.set_xlim(xmin, xmax)
            a.set_ylim(ymin, ymax)
            
        self.canvas.draw()

    def _on_button_release(self, event):
        'the release mouse button callback in pan/zoom mode'
        self.canvas.mpl_disconnect(self._motion_notify)
        if not self._xypress: return
        self._xypress = None
        self._button_pressed = None
        self.canvas.draw()
        self.plot.emit('pan-changed', self)


class ZoomPlotMode(PlotMode):
    def __init__(self, plot):
        PlotMode.__init__(self, plot, 'zoom',
                          'Zoom to rectangle','zoom_to_rect')        
        self._selectors = {}

    def activate(self):
        self.cids = []
        for ax in self.canvas.figure.get_axes():
            selector = Selector(ax, self._on_select,
                             select_type='rectangle')
            selector.line.set_linestyle('-')
            selector.line.set_linewidth(0.5)
            self.cids.append(self.canvas.mpl_connect('button_release_event',
                                                     selector._onrelease))
            self.cids.append(self.canvas.mpl_connect('motion_notify_event',
                                                     selector._onmove))
            self.cids.append(self.canvas.mpl_connect('button_press_event',
                                                     selector._onpress))
    def deactivate(self):
        for cid in self.cids:
            self.canvas.mpl_disconnect(cid)
    
    def _on_select(self, x1, y1, x2, y2, key):
        if abs(x1-x2)<0.0001 or abs(y1-y2)<0.0001:
            # too sensitive zoom crashes matplotlib
            return
        ax = self.canvas.figure.get_axes()[0]
        ax.set_xlim((min(x1, x2), max(x1, x2)))
        ax.set_ylim((min(y1,y2), max(y1, y2)))
        self.canvas.draw()
        self.plot.emit('zoom-changed', self)


class SelectLassoPlotMode(PlotMode):
    def __init__(self, plot):
        PlotMode.__init__(self, plot, 'lassoselect',
                          'Select within lasso', 'lasso')

        self._selectors = {}
        
    def activate(self):
        self.cids = []
        for ax in self.canvas.figure.get_axes():
            lasso = Selector(ax, self.plot.lasso_select_callback,
                             select_type='lasso')
            self.cids.append(self.canvas.mpl_connect('button_release_event',
                                                     lasso._onrelease))
            self.cids.append(self.canvas.mpl_connect('motion_notify_event',
                                                     lasso._onmove))
            self.cids.append(self.canvas.mpl_connect('button_press_event',
                                                     lasso._onpress))
    def deactivate(self):
        for cid in self.cids:
            self.canvas.mpl_disconnect(cid)


class SelectRectanglePlotMode(PlotMode):
    def __init__(self, plot):
        PlotMode.__init__(self, plot, 'select',
                          'Select within rectangle', 'select')
        self._selectors = {}
        
    def activate(self):
        self.cids = []
        for ax in self.canvas.figure.get_axes():
            lasso = Selector(ax, self.plot.rectangle_select_callback,
                             select_type='rectangle')
            self.cids.append(self.canvas.mpl_connect('button_release_event',
                                                     lasso._onrelease))
            self.cids.append(self.canvas.mpl_connect('motion_notify_event',
                                                     lasso._onmove))
            self.cids.append(self.canvas.mpl_connect('button_press_event',
                                                     lasso._onpress))
    def deactivate(self):
        for cid in self.cids:
            self.canvas.mpl_disconnect(cid)


class PlotToolbar(gtk.Toolbar):
    
    def __init__(self, plot):
        gtk.Toolbar.__init__(self)
        self.plot = plot
        self.canvas = plot.canvas
        self._current_mode = None
        self.tooltips = gtk.Tooltips()
        self._mode_sensitive = {}
        ## Maps toolbar buttons to PlotMode objects.
        self._mode_buttons = {}
        self.set_property('show-arrow', False)
        self.canvas.connect('enter-notify-event', self.on_enter_notify)
        self.show()
        
        self.add_mode(DefaultPlotMode(self.plot))
        self.add_mode(PanPlotMode(self.plot))
        self.add_mode(ZoomPlotMode(self.plot))
        self.add_mode(SelectRectanglePlotMode(self.plot))
        self.add_mode(SelectLassoPlotMode(self.plot))

        self.insert(gtk.SeparatorToolItem(), -1)
        self.set_style(gtk.TOOLBAR_ICONS)

         # Set up freeze button
        btn = gtk.ToolButton(gtk.STOCK_SAVE)
        
        #btn.set_icon_widget(image)
        btn.connect('clicked', self._on_save_clicked)
        self.insert(btn, -1)
        self.save_button = btn
        
        # Set up freeze button
        btn = gtk.ToggleToolButton()

        image = gtk.Image()
        image.set_from_pixbuf(fluents.icon_factory.get('freeze'))

        btn.set_icon_widget(image)
        btn.connect('toggled', self._on_freeze_toggle)
        self.insert(btn, -1)
        self.freeze_button = btn

        self.show_all()

    def add_mode(self, mode):
        """Adds a new mode to the toolbar."""

        self._mode_sensitive[mode] = True
        if len(self._mode_buttons) > 0:
            other = self._mode_buttons.keys()[0]
        else:
            other = None

        btn = gtk.RadioToolButton(other)
        btn.set_icon_widget(mode.get_icon())
        btn.set_tooltip(self.tooltips, mode.tooltip, 'Private')
        btn.connect('toggled', self._on_mode_toggle)

        self._mode_buttons[btn] = mode
        self.insert(btn, -1)

        if self._current_mode == None:
            self._current_mode = mode

    def get_mode(self):
        """Returns the active mode name."""
        if self._current_mode:
            return self._current_mode.name
        return None
    
    def get_mode_by_name(self, mode_name):
        """Returns the mode with the given name or None."""
        for m in self._mode_buttons.values():
            if m.name == mode_name:
                return m
        return None

    def get_button(self, mode_name):
        """Returns the button that corresponds to a mode name."""
        for b, m in self._mode_buttons.items():
            if m.name == mode_name:
                return b
        return None
    
    def set_mode(self, mode_name):
        """Sets a mode by name. Returns the mode or None"""
        if mode_name == self._current_mode.name:
            return None
        
        if self._current_mode:
            self._current_mode.deactivate()

        new_mode = self.get_mode_by_name(mode_name)
        if new_mode:
            new_mode.activate()
            self._current_mode = self.get_mode_by_name(mode_name)
        else:
            logger.log('warning', 'No such mode: %s' % mode_name)

        if self.get_button(mode_name) and \
               not self.get_button(mode_name).get_active():
            self.get_button(mode_name).set_active(True)
        globals()['active_mode'] = mode_name
        return self._current_mode

    def set_mode_sensitive(self, mode, bool):
        self._mode_sensitive[mode] = bool
        self.get_button(mode).set_sensitive(False)
        
    def _on_mode_toggle(self, button):
        mode = self._mode_buttons[button].name
        if self._mode_sensitive.get(mode) == False:
            logger.log('notice', 'Mode: %s is not sensitive in this plot' %mode)
        if button.get_active():
            self.set_mode(self._mode_buttons[button].name)

    def _on_save_clicked(self, button):
        """Pops up a file dialog saver, and saves current plot"""
        dialog = gtk.FileChooserDialog('Save plot')
        dialog.set_action(gtk.FILE_CHOOSER_ACTION_SAVE)
        dialog.add_buttons(gtk.STOCK_CANCEL, gtk.RESPONSE_CANCEL, gtk.STOCK_SAVE, gtk.RESPONSE_OK)
        default_name = "fluents_image"
        dialog.set_current_name("%s.png" % default_name)
        retval = dialog.run()
        if retval in [gtk.RESPONSE_CANCEL, gtk.RESPONSE_DELETE_EVENT]:
            logger.log("debug", "Cancelled save plot")
        elif retval == gtk.RESPONSE_OK:
            logger.log("debug", "Saving plot as: %s" % dialog.get_filename())
            self.canvas.print_figure(dialog.get_filename())
        else:
            print "unknown; ", retval 
        dialog.destroy()
        
    def _on_freeze_toggle(self, button):
        self.plot.set_frozen(button.get_active())

    def on_enter_notify(self, widget, event):
        self.set_mode(active_mode)
        # need views (plots) to grab key-events
        widget.grab_focus()


class Selector:
    def __init__(self, ax, callback, select_type='lasso'):
        self.drawon = True
        self.eventson = True
        self.select_type = select_type
        self.eventpress = None
        self.eventrelease = None
        self._background = None
        self.callback = callback
        self.axes = ax
        self.figure = self.axes.figure
        self.canvas = self.figure.canvas
        self.verts = []
        # to draw
        patch_props = dict(facecolor='gray', linewidth = 0,
                                 alpha=0.3, fill=True)
        line_props  = dict(linestyle='--', lw=1.0, color='black')

        self.line = Line2D([0,0], [0,0], visible=False, **line_props)
        self.axes.add_line(self.line)
        if self.select_type == 'lasso':
            self.polygon = Polygon([(0, 0)], visible=False, **patch_props)
        elif self.select_type == 'rectangle':
            self.polygon = Rectangle((0,0), 0, 1, visible=False, **patch_props)
        self.axes.add_patch(self.polygon)
    
    def _onrelease(self, event):
        if self.eventpress is None or self._ignore(event):
            return
        
        self.eventrelease = event
        self._show_selector(False)
        if self.verts is not None:
            self.verts.append((event.xdata, event.ydata))
            if self.select_type == 'lasso':
                if len(self.verts)>2:
                    self.callback(self.verts, event.key)
            elif self.select_type == 'rectangle':
                    x1, y1 = self.eventpress.xdata, self.eventpress.ydata
                    x2, y2 = self.eventrelease.xdata, self.eventrelease.ydata
                    self.callback(x1, y1, x2, y2, event.key)

        self.verts = []
        self.eventpress = None
        self.eventrelease = None
        self._background = None
        return False
        
    def _onmove(self, event):
        if self.eventpress is None or self._ignore(event):return

        # add new data
        if self.select_type == 'lasso':
            self.verts.append((event.xdata, event.ydata))
            self.polygon.xy = self.verts
            self.line.set_data(zip(*self.verts))
        
        elif self.select_type == 'rectangle':
            x,y = event.xdata, event.ydata
            minx, maxx = self.eventpress.xdata, x 
            miny, maxy = self.eventpress.ydata, y 
            if minx>maxx: minx, maxx = maxx, minx 
            if miny>maxy: miny, maxy = maxy, miny
            self.polygon.xy[0] = minx             
            self.polygon.xy[1] = miny
            self.polygon.set_width(maxx-minx)     
            self.polygon.set_height(maxy-miny)
            xx =  [minx, maxx, maxx, minx, minx]
            yy = [miny, miny, maxy, maxy, miny]
            self.line.set_data(xx,yy)
            
        # draw
        self.canvas.restore_region(self._background)
        self.axes.draw_artist(self.line)
        self.axes.draw_artist(self.polygon)
        self.canvas.blit(self.axes.bbox)

    def _onpress(self, event):
        # Is the event to be ignored?
        if self._ignore(event): return
        self._show_selector(True)
        self.eventpress = event
        self.verts = [(event.xdata, event.ydata)]
        return False

    def _ignore(self, event):
        'return True if event should be ignored'
        ignore = False
        # If no button was pressed yet ignore the event if it was out
        # of the axes
        if self.eventpress == None:
            return event.inaxes!= self.axes

        # If a button was pressed, check if the release-button is the
        # same.
        ignore = event.inaxes!=self.axes or\
                 event.button != self.eventpress.button
        
        if ignore:
            # remove poly + line if they are visible
            if self.polygon._visible or self.line._visible:
                self._show_selector(False)
        else:
           if not self.polygon._visible or not self.line._visible: 
               # put poly + line back if we event is ok
               self._show_selector(True)
    
        return ignore

    def _show_selector(self, bool):
        if self.polygon and self.line:
            self.polygon.set_visible(bool)
            self.line.set_visible(bool)
            if self._background == None:
                self._background = self.canvas.copy_from_bbox(self.axes.bbox)
            self.canvas.restore_region(self._background)
            self.canvas.blit()


# Create a view-changed signal that should be emitted every time
# the active view changes.
gobject.signal_new('view-changed', MainView, gobject.SIGNAL_RUN_LAST,
                   gobject.TYPE_NONE, 
                   (gobject.TYPE_PYOBJECT,))

# Create focus-changed signal
gobject.signal_new('focus-changed', ViewFrame, gobject.SIGNAL_RUN_LAST,
                   gobject.TYPE_NONE, 
                   (gobject.TYPE_PYOBJECT, gobject.TYPE_BOOLEAN,))

# Create zoom-changed signal
gobject.signal_new('zoom-changed', Plot, gobject.SIGNAL_RUN_LAST, None,
                   (gobject.TYPE_PYOBJECT,))

# Create pan/zoom-changed signal
gobject.signal_new('pan-changed', Plot, gobject.SIGNAL_RUN_LAST, None,
                   (gobject.TYPE_PYOBJECT,))

# Create plot-resize-changed signal
gobject.signal_new('plot-resize-changed', Plot, gobject.SIGNAL_RUN_LAST, None,
                   (gobject.TYPE_PYOBJECT,))