import pygtk
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
from system import logger


class MainView (gtk.Notebook):
    def __init__(self, view_listener):
        gtk.Notebook.__init__(self)
        self.set_show_tabs(False)
        self.set_show_border(False)

        # Add a multiple pane view and a single pane view.
        self.small_view = SmallView()
        self.small_view.set_view_listener(view_listener)
        self.small_view.show()
        self.large_view = LargeView()
        self.large_view.show()

        self.append_page(self.small_view)
        self.append_page(self.large_view)
        self.set_current_page(0)
    
    def goto_small(self):
        if self.get_current_page() == 0:
            return None
        self.set_current_page(0)
        view = self.large_view.remove_child()
        self.small_view.return_current(view)

    def goto_large(self):
        if self.get_current_page() == 1:
            return None
        self.set_current_page(1)
        view = self.small_view.borrow_current()
        self.large_view.set_child(view)

    def show(self):
        gtk.Notebook.show(self)

    def insert_view(self, view):
        self.small_view.insert_view(view)


class Plot (gtk.Frame):

    def __init__(self, title):
        gtk.Frame.__init__(self)
        self.mark_active(False)
        self.connect('button_press_event', self.on_button_press)
        self.sel_obj = None
        self.active = False
        self.title = title
        self.selection_listener = None

    def get_title(self):
        return self.title

    def on_button_press(self, *rest):
#        logger.log('debug', 'button pressed in plot')
        self.mark_active(True)

    def mark_active(self, active):
        if active:
            self.set_shadow_type(gtk.SHADOW_IN)
        else:
            self.set_shadow_type(gtk.SHADOW_OUT)
        self.active = active

    def selection_changed(self, selection):
        pass

    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 SmallView (gtk.Table):
        
    def __init__(self):
        gtk.Table.__init__(self, 2, 2, True)

        self.child_views = [[EmptyView(), EmptyView()],
                            [EmptyView(), EmptyView()]]
        self.cols = 2
        self.rows = 2

        self.active_x = 0
        self.active_y = 0

        self.set_row_spacings(3)
        self.set_col_spacings(3)
        self._listener = None
        
        for x in range(self.cols):
            for y in range(self.rows):
                child = self.child_views[x][y]
                child.parent_signalling = child.connect('button_press_event', self.__view_button_event__)
                self.attach(child, x, x+1, y, y+1)

    def set_view_listener(self, listener):
        self._listener = listener

    def set_child(self, child, col, row):
        cur_widget = self.child_views[col][row]
        cur_widget.disconnect(cur_widget.parent_signalling)
        self.remove(cur_widget)
        self.attach(child, col, col+1, row, row+1)
        child.parent_signalling = child.connect('button_press_event', self.__view_button_event__)
        self.child_views[col][row] = child
        if cur_widget.active:
            child.mark_active(True)
            cur_widget.mark_active(False)
        child.show()

    def borrow_current(self):
        self.borrowed = self.child_views[self.active_x][self.active_y]
        self.remove_child(self.active_x, self.active_y)
        return self.borrowed

    def return_current(self, view):
        self.set_child(view, self.active_x, self.active_y)

    def remove_child(self, col, row):
        self.remove(self.child_views[col][row])
        self.attach(EmptyView(), col, col+1, row, row+1)
       
    def insert_view(self, child):
        if not self.find_child(child):
            self.set_child(child, self.active_x, self.active_y)

    def show(self):
        for x in self.child_views:
            for y in x:
                y.show()
        gtk.Table.show(self)

    def hide(self):
        for x in self.child_views:
            for y in x:
                y.hide()
        gtk.Table.hide(self)

    def set_active(self, x, y):
        old_focus = self.child_views[self.active_x][self.active_y]
        new_focus = self.child_views[x][y]
        old_focus.mark_active(False)
        new_focus.mark_active(True)
        self.active_x = x
        self.active_y = y
        if self._listener:
            logger.log("debug", "emitting")
            self._listener(new_focus)

    def find_child(self, child):
        for i, row in enumerate(self.child_views):
            for j, v in enumerate(row):
                if v == child:
                    return (i, j)
        return None

    def __view_button_event__(self, view, *rest):
        loc = self.find_child(view)
        if loc:
            self.set_active(loc[0], loc[1])

    def get_view(self, x, y):
        return self.child_views[x][y]

class LargeView (gtk.Frame):
    def __init__(self):
        gtk.Frame.__init__(self)
        self.child_view = EmptyView()
        self.add(self.child_view)
    
    def set_child(self, child):
        self.remove(self.child_view)
        self.child_view.hide()
        self.add(child)
        self.child_view = child
        child.show()

    def hide(self):
        self.child_view.hide()
        gtk.Frame.hide(self)

    def show(self):
        self.child_view.show()
        gtk.Frame.show(self)

    def remove_child(self):
        child = self.child_view
        child.hide()
        self.remove(child)
        return child
        
    
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.connect('button_press_event', self.on_button_press)

        self.label = label
        self.ebox = ebox
        self.show()

    def show(self):
        self.ebox.show()
        self.label.show()
        Plot.show(self)
        
    def hide(self):
        self.label.hide()
        self.ebox.hide()
        Plot.hide(self)


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

    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.draw()
        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 SinePlot(Plot):
    def __init__(self):
        Plot.__init__(self, 'Sine plot')

        fig = Figure(figsize=(5,4), dpi=72)
        ax = fig.add_subplot(111)
        t = arange(0.0,3.0,0.01)
        s = sin(2*pi*t)
        ax.plot(t,s)

        self.canvas = FigureCanvas(fig)
        self._toolbar = NavToolbar(self.canvas, None)
        self._toolbar.set_property('show-arrow', False)
        self.add(self.canvas)
        self.canvas.show()

    def get_toolbar(self):
        return self._toolbar


class LinePlot(Plot):
    def __init__(self, dataset, name="Line plot"):
        Plot.__init__(self, name)
        fig = Figure(figsize=(5,4), dpi=72)
        self.canvas = FigureCanvas(fig)

        self.ax = ax = fig.add_subplot(111)
        self._dataset = dataset

        import rpy
        silent_eval = rpy.with_mode(rpy.NO_CONVERSION, rpy.r)

        rpy.with_mode(rpy.NO_CONVERSION, rpy.r.assign)("m", dataset.get_matrix())        
        ymin, ymax = rpy.r("range(m)*1.1")
        self._ymin, self._ymax = ymin, ymax
        silent_eval("res <- apply(m, 2, density, from=%s, to=%s)" % (ymin, ymax))
        silent_eval("k <- sapply(1:length(res), function(id) rev(res[[id]]$y))")
        self._bg_matrix = bg_matrix = rpy.r("k")
        rpy.r.rm(["k", "res", "m"])

        # Hack - we draw plot in selection_changed()
        self.selection_changed(None)

        self.add(self.canvas)
        self.canvas.show()

        # We use a regular toolbar as we don't need selections
        self._toolbar = NavigationToolbar2(self.canvas, None)
        self._toolbar.set_property('show-arrow', False)

    def get_toolbar(self):
        self.canvas.draw()
        return self._toolbar

    def selection_changed(self, selection):
        self.ax.clear()
        
        rows, cols = self._bg_matrix.shape
        self.ax.imshow(self._bg_matrix, cmap=cm.Greys, extent=(0.5, cols+0.5, self._ymin, self._ymax))

        dim_2, dim_1 = self._dataset.get_dim_names()

        if selection:
            ids = selection[dim_2] # current identifiers
            index = [ind for id,ind in self._dataset[dim_2].items() if id in ids] #conversion to index
            for i in index:
                line = self._dataset.get_matrix()[i]
                self.ax.plot(range(1, len(line)+1), line)
            
        self.ax.set_xlim((0.5, cols+0.5))
        self.ax.set_ylim((self._ymin, self._ymax))
        self.canvas.draw()


class ScatterPlot(Plot):
    def __init__(self, dataset, id_dim, sel_dim, id_1, id_2, name="Scatter plot"):
        Plot.__init__(self, name)
        fig = Figure(figsize=(5,4), dpi=72)
        self.ax = ax = fig.add_subplot(111)
        self.current_dim = id_dim
        # testing testing
        self.dataset = dataset
        x_index = dataset[sel_dim][id_1]
        y_index = dataset[sel_dim][id_2]
        
        self.xaxis_data = dataset._array[:,x_index]
        self.yaxis_data = dataset._array[:,y_index]
        ax.plot(self.xaxis_data,self.yaxis_data,'og')
        ax.set_title(self.get_title())
        ax.set_xlabel("%s - %s" % (sel_dim, id_1))
        ax.set_ylabel("%s - %s" % (sel_dim, id_2))

        ###

        self.canvas = FigureCanvas(fig)
        self.add(self.canvas)
        self.canvas.show()

        self._toolbar = NavToolbar(self.canvas, None)
        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):
        'event1 and event2 are the press and release events'
        ydata = self.yaxis_data
        xdata = self.xaxis_data

        # find indices of selected area
        if x1>x2:
            if y1<y2:
                index =scipy.nonzero((xdata<x1) & (xdata>x2) & (ydata>y1) & (ydata<y2))
            else:
                index =scipy.nonzero((xdata<x1) & (xdata>x2) & (ydata<y1) & (ydata>y2))
        else:
            #logger.log('debug','Selection x_start less than x_end')
            if y1<y2:
                #logger.log('debug','Selection y_start less than y_end')
                index =scipy.nonzero((xdata>x1) & (xdata<x2) & (ydata>y1) & (ydata<y2))
            else:
                #logger.log('debug','Selection y_start bigger than y_end')
                index =scipy.nonzero((xdata>x1) & (xdata<x2) & (ydata<y1) & (ydata>y2))

        # generate ids for selected indices
        reverse = {}
        for key, value in self.dataset[self.current_dim].items():
            reverse[value] = key
        ids = []
        for ind in index:
            ids.append(reverse[ind])
            
        self.selection_listener(self.current_dim,ids)

    def selection_changed(self, selection):
        ids = selection[self.current_dim] # current identifiers
        
        index = [ind for id,ind in self.dataset[self.current_dim].items() if id in ids] #conversion to index
        xdata_new = scipy.take(self.xaxis_data,index) #take data
        ydata_new = scipy.take(self.yaxis_data,index)
        self.ax.clear()
        self.ax.plot(self.xaxis_data,self.yaxis_data,'og')
        self.ax.plot(xdata_new,ydata_new,'or')
        self.canvas.draw()