import os,sys
import pygtk
import gobject
import gtk
import matplotlib
import scipy
from matplotlib.backends.backend_gtkagg import FigureCanvasGTKAgg as FigureCanvas
from matplotlib.backends.backend_gtkagg import NavigationToolbar2GTK as NavigationToolbar2
from matplotlib.axes import Subplot
from matplotlib.figure import Figure
from matplotlib.numerix import arange, sin, pi
from matplotlib.widgets import RectangleSelector
from matplotlib import cm
import networkx
from system import logger


class ObjectTable:
    """A 2D table of elements."""

    def __init__(self, xsize=0, ysize=0, creator=None):
        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."""
        # 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)])

    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 bin 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):
        gtk.Frame.__init__(self)
        self.focused = False
        self.view_frames = view_frames
        self.empty_view = EmptyView()
        self._button_event = None
        
        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.add(self._view)
        self._view.show()
        self.show()

    def focus(self):
        """Gets focus and ensures that no other window is in focus."""
        if self.focused:
            return self

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

        self.set_shadow_type(gtk.SHADOW_IN)
        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.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"""

        # 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
        view_parent = view.get_parent()
        if view_parent:
            view_parent.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
        self._view.hide()
        self.remove(self._view)
        self.add(view)
        view.show()
        
        self._view = view

    def get_view(self):
        """Returns current view, or None if the empty view is set."""
        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)


class MainView (gtk.Notebook):
    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._small_views = gtk.Table(2, 2, True)
        self._small_views.set_col_spacings(4)
        self._small_views.set_row_spacings(4)
        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):
        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)

    def get_active_small_view(self):
        for vf in self._view_frames:
            if vf.focused:
                return vf
        return None
    
    def goto_large(self):
        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):
        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):
        if self.get_current_page() == 0:
            vf = self.get_active_small_view()
        else:
            vf = self._large_view
        vf.set_view(view)
        
    def show(self):
        for vf in self._view_frames:
            vf.show()
        self._small_views.show()
        gtk.Notebook.show(self)
        
    def on_view_focus_changed(self, widget, vf, focused):
        if focused:
            self.emit('view-changed', vf)
    

class Plot (gtk.Frame):
    def __init__(self, title):
        gtk.Frame.__init__(self)
        self.sel_obj = None
        self.title = title
        self.selection_listener = None
        self.set_shadow_type(gtk.SHADOW_NONE)
        self._background = None
        self.canvas = None
        self.fig = Figure(figsize=(5,4), dpi=72)
        self.canvas = FigureCanvas(self.fig)
        self.fig.set_facecolor('white')
        self._sel_sensitive = True

    def set_selection_sensitive(self,event):
        if event:
            if event.get_active():
                logger.log('debug','Selection freezed')
                self._sel_sensitive = False
            else:
                logger.log('debug','Selections active')
                self._sel_sensitive = True

    def get_title(self):
        return self.title

    def selection_changed(self, selection):
        if not self._sel_sensitive or not self.get_property('visible'):
            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 get_toolbar(self):
        return None


class EmptyView (Plot):
    def __init__(self):
        Plot.__init__(self, 'Empty view')
        label = gtk.Label('No view')
        ebox = gtk.EventBox()
        ebox.add(label)
        self.add(ebox)
        
        label.show()
        ebox.show()
        self.show()


class NavToolbar(NavigationToolbar2):
    toolitems = (('Select', 'Select within rectangle', 'zoom_to_rect.png',
                  'select'),) + NavigationToolbar2.toolitems
    def __init__(self, *args):
        NavigationToolbar2.__init__(self, *args)
        self._select_callback = None
        self.message.hide()

    def select(self, *args):
        """Selection mode selected handler."""
        if self._active == 'SELECT':
            self._active = None
        else:
            self._active = 'SELECT'

        if self._idPress is not None:
            self._idPress = self.canvas.mpl_disconnect(self._idPress)
            self.mode = ''

        if self._idRelease is not None:
            self._idRelease = self.canvas.mpl_disconnect(self._idRelease)
            self.mode = ''

        if self._active:
            self._idPress = self.canvas.mpl_connect('button_press_event', self.press_select)
            self._idRelease = self.canvas.mpl_connect('button_release_event', self.release_select)
            self.mode = 'Select rectangle mode'
        self.set_message(self.mode)

    def set_message(self, s):
        """Set status in toolbar to string s.

        Overrided to make sure message can be updated even when
        drawing rubberband.
        """
        self.message.set_label(s)

    def press_select(self, event):
        """Mouse button pressed handler for selection mode."""
        x, y = event.x, event.y

        for i, a in enumerate(self.canvas.figure.get_axes()):
            if event.inaxes==a and event.inaxes.get_navigate():
                xmin, xmax = a.get_xlim()
                ymin, ymax = a.get_ylim()
                lim = xmin, xmax, ymin, ymax
                self._xypress = x, y, a, i, lim, a.transData.deepcopy()
                break
        
        self.press(event)

    def release_select(self, event):
        """Mouse button released handler for selection mode."""
        # only release if button was pressed inside first?
        if self._xypress:
            x, y = event.x, event.y
            lastx, lasty, a, ind, lim, trans = self._xypress
            lastx, lasty = a.transData.inverse_xy_tup( (lastx, lasty) )
            x, y = a.transData.inverse_xy_tup( (x, y) )

            if self._select_callback:
                self._select_callback(lastx, lasty, x, y)
        
        self._xypress = None
        self.push_current()
        self.release(event)

    def mouse_move(self, event):
        """Extend NavigationToolbar2.mouse_move to provide selection support."""
        from matplotlib.backend_bases import cursors

        # Only update the rubberband for selection mode when mouse is
        # within the plotting area.
        if event.inaxes and self._active=='SELECT':
            if self._lastCursor != cursors.SELECT_REGION:
                self.set_cursor(cursors.SELECT_REGION)
                self._lastCursor = cursors.SELECT_REGION
            if self._xypress is not None:
                x, y = event.x, event.y
                lastx, lasty, a, ind, lim, trans= self._xypress
                self.draw_rubberband(event, x, y, lastx, lasty)

        NavigationToolbar2.mouse_move(self, event)

    def set_select_callback(self, listener):
        """Allow plots to register a callback for selection events.

        The callback will be called as listener(x1, y1, x2, y2).  All
        coordinates are in the plot coordinate system, not pixels or
        widget coordinates.
        """
        self._select_callback = listener


class PlotToolbar(NavToolbar):
    """Extensions to existing toolbar
    """
    def __init__(self, *args):
        NavToolbar.__init__(self, *args)
        self.tb_freeze = gtk.ToolItem()
        self.chk = gtk.CheckButton ()
        self.chk.set_label ('Freeze')
        self.tb_freeze.add (self.chk)
        self.tb_freeze.set_tooltip(self.tooltips, 'Freeze current selection')
        self.insert(self.tb_freeze,-1)        
        toolitem = gtk.SeparatorToolItem()
        self.insert(toolitem, -1)
        self.show_all()
        self.message.hide()


class LineViewPlot(Plot):
    """Line view of current selection, no interaction
    Only works on 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
    ps: slow (cant get linecollection and blit to work)
    """
    def __init__(self, dataset,major_axis=1,minor_axis=None, name="Line view"):
        self.use_blit = False
        self._last_index = []
        self._data = dataset.asarray()
        self.dataset = dataset
        Plot.__init__(self, name)
        self.ax = self.fig.add_subplot(111)
        self.ax.set_title(self.get_title())
        self.current_dim = self.dataset.get_dim_name(major_axis)
        if len(self._data.shape)==2 and not minor_axis:
            minor_axis = major_axis-1

        #initial draw
        self.line_collection = {}
        x_axis = scipy.arrayrange(self._data.shape[minor_axis])
        for i in range(self._data.shape[major_axis]):
            yi = scipy.take(self._data,[i],axis=major_axis)
            if self.use_blit:
                l,=self.ax.plot(x_axis,yi,'k',alpha=.05,animated=True)
            else:
                l,=self.ax.plot(x_axis,yi,'k',alpha=.05)
            self.line_collection[i] = l
        
        self.add(self.canvas)
        self.canvas.show()
        
        #FIXME: Lineview plot cannot do selections -> disable in toolbar
        self._toolbar = PlotToolbar(self.canvas, None)
        self._toolbar.set_property('show-arrow', False)
        self._toolbar.chk.connect ('toggled' ,self.set_selection_sensitive)
        self.canvas.mpl_connect('resize_event',self.clear_background)


    def get_toolbar(self):
        return self._toolbar
        
    def clear_background(self,event):
        self._background = None
        
    def set_current_selection(self, selection):
        ids = selection[self.current_dim] # current identifiers
        index = self.dataset.get_indices(self.current_dim, ids)
        if self.use_blit:
            if self._background is None:
                self._last_index = None
                self._background = self.canvas.copy_from_bbox(self.ax.bbox)
            self.canvas.restore_region(self._background)

        if index:
            if self._last_index:
                for i in self._last_index:
                    # if not using blit: reset last selection
                    self.ax.lines[i].set_color('k')
                    self.ax.lines[i].set_alpha(.05)
                    self.ax.lines[i].set_zorder(1)
            for i in index:
                self.ax.lines[i].set_color('r')
                self.ax.lines[i].set_alpha(1.0)
                self.ax.lines[i].set_visible(True)
                self.ax.lines[i].set_zorder(3)
                if self.use_blit:
                    self.ax.draw_artist(self.ax.lines[i])
                    
                self._last_index = index

            if self.use_blit:    
                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.ax = self.fig.add_subplot(111)
        self.ax.axhline(0,color='k',lw=1.5,zorder=0)
        self.ax.axvline(0,color='k',lw=1.5,zorder=0)
        self.current_dim = id_dim
        self.dataset_1 = dataset_1
        self.ms = s
        self._selection_line = None
        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]
        self.ax.plot(self.xaxis_data,self.yaxis_data,'o',markeredgewidth=0,markersize=s)
        self.ax.set_title(self.get_title())
        self.add(self.canvas)
        self.canvas.show()

        # add toolbar
        self._toolbar = PlotToolbar(self.canvas, None)
        self._toolbar.chk.connect ('toggled' ,self.set_selection_sensitive)
        self._toolbar.set_property('show-arrow', False)
        self._toolbar.set_select_callback(self.rectangle_select_callback)

    def get_toolbar(self):
        return self._toolbar

    def rectangle_select_callback(self, x1, y1, x2, y2):
        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))
        ids = self.dataset_1.get_identifiers(self.current_dim, index)
        self.selection_listener(self.current_dim, ids)

    def set_current_selection(self, selection):
        ids = selection[self.current_dim] # current identifiers
        index = self.dataset_1.get_indices(self.current_dim, ids)
        xdata_new = scipy.take(self.xaxis_data, index) #take data
        ydata_new = scipy.take(self.yaxis_data, index)
        #remove old selection
        if self._selection_line:
            self.ax.lines.remove(self._selection_line)

        self._selection_line, = self.ax.plot(xdata_new,ydata_new,marker='o',markersize=self.ms,linestyle=None,markerfacecolor='r')

        self._toolbar.forward() #update data lims before draw
        self.canvas.draw()


class ScatterPlot(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,c='b',s=30, name="Scatter plot"):
        Plot.__init__(self, name)
        self.ax = self.fig.add_subplot(111)
        self.current_dim = id_dim
        self.dataset_1 = dataset_1
        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]
        lw = scipy.zeros(self.xaxis_data.shape,'f')
        self.ax.scatter(self.xaxis_data,self.yaxis_data,s=s,c=c,linewidth=lw,edgecolor='k',alpha=.6,cmap = cm.Set1)
        self.ax.set_title(self.get_title())
        # collection
        self.coll = ax.collections[0]

        # add canvas to widget
        self.add(self.canvas)
        self.canvas.show()

        # create toolbar
        self._toolbar = PlotToolbar(self.canvas, None)
        self._toolbar.chk.connect ('toggled' ,self.set_selection_sensitive)
        self._toolbar.set_select_callback(self.rectangle_select_callback)

    def get_toolbar(self):
        return self._toolbar

    def rectangle_select_callback(self, x1, y1, x2, y2):
        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))

        ids = self.dataset_1.get_identifiers(self.current_dim, index)
        self.selection_listener(self.current_dim, ids)

    def set_current_selection(self, selection):
        ids = selection[self.current_dim] # current identifiers
        index = self.dataset_1.get_indices(self.current_dim, ids)
        lw = scipy.zeros(self.xaxis_data.shape,'f')
        scipy.put(lw,index,2.)
        zo = lw.copy() + 1 #z-order, selected on top
        self.coll.set_linewidth(lw)
        self.coll.set_zorder(zo)
        self._toolbar.forward() #update data lims before draw
        self.canvas.draw()

        
class NetworkPlot(Plot):
    def __init__(self, dataset, **kw):
        # Set member variables and call superclass' constructor
        self.graph = dataset.asnetworkx()
        self.dataset = dataset 
        self.keywords = kw
        self.dim_name = self.dataset.get_dim_name(0)
        if not kw.has_key('name'):
            kw['name'] = self.dataset.get_name()
        if not kw.has_key('prog'):
            kw['prog'] = 'neato'
        if not kw.has_key('pos') or kw['pos']:
            kw['pos'] = networkx.drawing.nx_pydot.graphviz_layout(self.graph, kw['prog'])
        Plot.__init__(self, kw['name'])

        # Keep node size and color as dicts for fast lookup
        self.node_size = {}
        if kw.has_key('node_size') and cb.iterable(kw['node_size']):
            kw.remove('node_size')
            for id, size in zip(self.dataset[self.dim_name], kw['node_size']):
                self.node_size[id] = size
        else:
            for id in dataset[self.dim_name]:
                self.node_size[id] = 40
                
        self.node_color = {}
        if kw.has_key('node_color') and cb.iterable(kw['node_color']):
            kw.remove('node_color')
            for id, color in zip(self.dataset[self.dim_name], kw['node_color']):
                self.node_color[id] = color
        else:
            self.node_color = None
#            for id in self.dataset[self.dim_name]:
#                self.node_color[id] = 'red'

        if kw.has_key('node_color'):
            kw.pop('node_color')

        self.ax = self.fig.add_subplot(111)
        self.ax.set_position([0.01,0.01,.99,.99])
        self.ax.set_xticks([])
        self.ax.set_yticks([])
        # FIXME: ax shouldn't be in kw at all
        if kw.has_key('ax'):
            kw.pop('ax')

        # Add canvas and show
        self.add(self.canvas)
        self.canvas.show()

        # Initial draw
        networkx.draw_networkx(self.graph, ax=self.ax, **kw)

        # Setup toolbar
        self._toolbar = PlotToolbar(self.canvas, None)
        self._toolbar.chk.connect ('toggled' ,self.set_selection_sensitive)
        self._toolbar.set_select_callback(self.rectangle_select_callback)

    def get_toolbar(self):
        return self._toolbar

    def rectangle_select_callback(self, x1, y1, x2, y2):
        'event1 and event2 are the press and release events'
        pos = self.keywords['pos']
        ydata = scipy.zeros((len(pos),),'l')
        xdata = scipy.zeros((len(pos),),'l')
        node_ids = []
        c = 0
        for name,(x,y) in pos.items():
            node_ids.append(name)
            xdata[c] = x
            ydata[c] = y
            c+=1
        
        # find indices of selected area
        if x1 > x2:
            x1, x2 = x2, x1
        if y1 > y2:
            y1, y2 = y2, y1
        index = scipy.nonzero((xdata>x1) & (xdata<x2) & (ydata>y1) & (ydata<y2))

        ids = [node_ids[i] for i in index]
        self.selection_listener(self.dataset.get_dim_name(0), ids)

    def set_current_selection(self, selection):
        ids = selection[self.dataset.get_dim_name(0)] # current identifiers
        node_set = set(self.graph.nodes())

        selected_nodes = list(ids.intersection(node_set))
        unselected_nodes = list(node_set.difference(ids))

        if self.node_color:
            unselected_colors = [self.node_color[x] for x in unselected_nodes]
        else:
            unselected_colors = 'red'

        if self.node_size:
            unselected_sizes = [self.node_size[x] for x in unselected_nodes]
            selected_sizes = [self.node_size[x] for x in selected_nodes]

        self.ax.clear()
        networkx.draw_networkx_edges(self.graph, edge_list=self.graph.edges(), \
            ax=self.ax, **self.keywords)
        networkx.draw_networkx_labels(self.graph,**self.keywords)
        if unselected_nodes:
            networkx.draw_networkx_nodes(self.graph, nodelist=unselected_nodes, \
                node_color='r', node_size=unselected_sizes, ax=self.ax, **self.keywords)

        if selected_nodes:
            networkx.draw_networkx_nodes(self.graph, nodelist=selected_nodes, \
             node_color='k', node_size=selected_sizes, ax=self.ax, **self.keywords)
        
        self.canvas.draw()


# 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,))