import pygtk
import gtk
import gtk.gdk
import gtk.glade
import gnome
import gnome.ui
import gobject

import logger, dataset
import annotations
from lib import hypergeom


class SimpleMenu(gtk.Menu):
    def __init__(self):
        gtk.Menu.__init__(self)
        
    def add_simple_item(self, title, function):
        item = gtk.MenuItem(title)
        item.connect('activate', function)
        self.append(item)
        item.show()


class IdListController:
    """Controller class for the identifier list."""

    def __init__(self, idlist):
        self._idlist = idlist

        # dimname: current_annotation_name
        self._annotation = {}

        # current dimension
        self._dimension = None

        # id, annotation
        self._idstore = gtk.ListStore(gobject.TYPE_STRING,
                                      gobject.TYPE_STRING,)
        
        # Annotation tree column
        self._annotation_column = None

        ## Set up identifier list
        idlist.set_model(self._idstore)
        
        renderer = gtk.CellRendererText()
        dim_column = gtk.TreeViewColumn('Identifiers', renderer, text=0)
        idlist.insert_column(dim_column, 0)
        idlist.connect('button-press-event', self._button_pressed)

        ## Set up identifier list context menu
        menu = self._menu = SimpleMenu()
        menu.add_simple_item('Import...', self._on_import_list)
        menu.add_simple_item('Export...', self._on_export_list)
        
        item = gtk.MenuItem('Show annotations')
        menu.append(item)
        item.show()
        self._menu_ann = item
       
    ##
    ## Public interface
    ##
    def set_dimension(self, dimname):
        """Set dimension"""
        if dimname == self._dimension:
            return

        self.set_annotation(self._annotation.get(dimname, None))

        if not self._annotation.has_key(dimname):
            self._annotation[dimname] = None
        self._dimension = dimname 

    def set_annotation(self, annotation):
        """Set the displayed annotation to annotation.  If annotation is None,
        the annotation column is hidden. Otherwise the annotation column is
        shown and filled with values from the given annotation field."""

        if annotation == None:
            if self._annotation_column != None:
                self._idlist.remove_column(self._annotation_column)
            self._annotation_column = None
        else:
 
            idlist = [x[0] for x in self._idstore]
            annlist = annotations.get_dim_annotations(self._dimension, 
                          annotation,
                          idlist)

            for i, x in enumerate(self._idstore):
                x[1] = annlist[i]

            if self._annotation_column == None:
                renderer = gtk.CellRendererText()
                col = gtk.TreeViewColumn(annotation, renderer, text=1)
                self._idlist.append_column(col)
                self._annotation_column = col
            self._annotation_column.set_title(annotation)

        self._annotation[self._dimension] = annotation

    def set_selection(self, selection):
        """Set the selection to be displayed.
        The selection is not stored, the values are copied into the TreeStore"""
        
        self._idstore.clear()
        id_list = list(selection[self._dimension])
        
        idlist = list(selection[self._dimension])
        if self._annotation[self._dimension] != None:
            annlist = annotations.get_dim_annotations(self._dimension, 
                          self._annotation[self._dimension],
                          idlist)
            for id, ann in zip(idlist, annlist):
                self._idstore.append((id, ann))
        else:
            for e in idlist:
                self._idstore.append((e, None))

    ##
    ## Private interface
    ##
    def _update_annotations_menu(self):
        """Updates the annotations menu with the available annotations for the
        current dim."""

        dim_h = annotations.get_dim_handler(self._dimension)
        if not dim_h:
            self._menu_ann.set_sensitive(False)
        else:
            annotations_menu = gtk.Menu()
            self._menu_ann.set_sensitive(True)
            dh = annotations.get_dim_handler(self._dimension)
            ann_names = dh.get_annotation_names()

            for ann in ann_names:
                item = gtk.MenuItem(ann)
                item.connect('activate', self._on_annotation_activated, ann)
                annotations_menu.append(item)
                item.show()

            self._menu_ann.set_submenu(annotations_menu)

    ##
    ## GTK Callbacks
    ##
    def _popup_menu(self, *rest):
        self._update_annotations_menu()
        self._menu.popup(None, None, None, 0, 0)

    def _on_annotation_activated(self, menuitem, annotation):
        self.set_annotation(annotation)

    def _button_pressed(self, widget, event):
        if event.button == 3:
            self._update_annotations_menu()
            self._menu.popup(None, None, None, event.button, event.time)
        
    def _on_export_list(self, menuitem):
        print "export stuff"

    def _on_import_list(self, menuitem):
        dialog = gtk.FileChooserDialog('Load annotations')
        dialog.set_action(gtk.FILE_CHOOSER_ACTION_OPEN)
        dialog.add_buttons(gtk.STOCK_CANCEL, gtk.RESPONSE_CANCEL, 
                           gtk.STOCK_OPEN, gtk.RESPONSE_OK)
        dialog.set_select_multiple(True)
        retval = dialog.run()
        if retval in [gtk.RESPONSE_CANCEL, gtk.RESPONSE_DELETE_EVENT]:
            pass
        elif retval == gtk.RESPONSE_OK:
            for filename in dialog.get_filenames():
                annotations.read_annotations_file(filename)
        else:
            print "unknown; ", retval 
        dialog.destroy()


class SelectionListController:
    def __init__(self, seltree, idlist_controller):
        self._seltree = seltree
        self._sel_stores = {}
        self._dimension = None
        self._idlist_controller = idlist_controller
    
        # Selection renderers
        renderer = gtk.CellRendererText()
        sel_column = gtk.TreeViewColumn('Selection', renderer, text=0)
        seltree.insert_column(sel_column, 0)
        
        intersect_cs_col = gtk.TreeViewColumn('In curr.sel.', renderer, text=3)
        seltree.insert_column(intersect_cs_col, 1)
    
        n_all_col = gtk.TreeViewColumn('# All', renderer, text=4)
        seltree.insert_column(n_all_col, 2)
        
        rank_col = gtk.TreeViewColumn('Rank', renderer, text=5)
        seltree.insert_column(rank_col, 3)

        # Signals
        seltree.connect('row-activated', self._on_row_activated)
        seltree.connect('cursor-changed', self._on_cursor_changed)
        seltree.connect('button-press-event', self._on_button_pressed)
        seltree.drag_dest_set(gtk.DEST_DEFAULT_ALL,
                              [("GTK_TREE_MODEL_ROW", gtk.TARGET_SAME_APP, 7)],
                              gtk.gdk.ACTION_LINK)
        
        seltree.connect('drag-data-received', self._drag_data_received)

        # Selections context menu
        self._seltree_menu = SimpleMenu()
        self._seltree_menu.add_simple_item('Sort by selection', self._on_seltree_sort)
        # self._seltree_menu.add_simple_item('Copy selection', self._on_seltree_copy_selection)
    #
    # Public interface
    #

    def set_project(self, project):
        """Dependency injection."""
        self.project = project
        project.add_selection_observer(self)
        
    def set_dimlist_controller(self, dimlist_controller):
        self._dimlist_controller = dimlist_controller

    def set_dimension(self, dim):
        self._ensure_selection_store(dim)
        self._seltree.set_model(self._sel_stores[dim])
        self._idlist_controller.set_dimension(dim)
        self._dimension = dim

    def selection_changed(self, dimname, selection):
        """Callback function from Project."""
        for dim in selection.dims():
            self._ensure_selection_store(dim)
            store = self._sel_stores[dim]

            if not self._get_current_selection_iter(selection, dim):
                n = len(selection[dim])
                values = (selection.title, selection, dim, n, n, 2)
                store.insert_after(None, None, values)

    def add_dataset(self, dataset):
        dim_name = dataset.get_dim_name(0)
        self._ensure_selection_store(dim_name)
        store = self._sel_stores[dim_name]
        di = self._get_dataset_iter(dataset)
        if not di:
            n_tot = dataset.shape[0]
            selection = self.project.get_selection().get(dim_name)
            ds_idents = dataset.get_identifiers(dim_name)
            n_cs = len(selection.intersection(ds_idents))
            values = (dataset.get_name(), dataset, dim_name, n_cs, n_tot, 2)

            i = store.insert_after(None, None, values)
            for selection in dataset.as_selections():
                n_sel = len(selection[dim_name])
                values = (selection.title, selection, dim_name, 0, n_sel, 0)
                store.insert_after(i, None, values)

        
    ##
    ## Private interface
    ##
    def _add_selection_store(self, dim):
        """Add a new gtk.TreeStore for the selections on a dimension."""
        # Create new store
        # Two types of lines, one for CategoryDatasets and one for
        # Selections. The elements are title, link to dataset or selection,
        # name of dimension, num. members in selection, num. in
        # intersection with current selection and the rank of selection.
        store = gtk.TreeStore(gobject.TYPE_STRING,
                              gobject.TYPE_PYOBJECT,
                              gobject.TYPE_STRING,
                              gobject.TYPE_INT,
                              gobject.TYPE_INT,
                              gobject.TYPE_FLOAT)

        # Set selection store for this dimension
        self._sel_stores[dim] = store

    def _ensure_selection_store(self, dim):
        """Ensure that the object has a gtk.TreeStore for the given dimension"""
        # Do not overwrite existing stores
        if self._sel_stores.has_key(dim):    
            return
        self._add_selection_store(dim)

    def _get_dataset_iter(self, ds):
        """Returns the iterator to the selection tree row containing a
        given dataset."""

        store = self._sel_stores[ds.get_dim_name(0)]

        i = store.get_iter_first()
        while i:
            if store.get_value(i, 1) == ds:
                return i
            i = store.iter_next(i)
        return None

    def _get_current_selection_iter(self, selection, dimension):
        if not self._sel_stores.has_key(dimension):
            return None

        store = self._sel_stores[dimension]

        i = store.get_iter_first()
        while i:
            if store.get_value(i, 1) == selection:
                if store.get_value(i, 2) == dimension:
                    return i
            i = store.iter_next(i)
        return None

    def _sort_selections(self, dataset):
        """Ranks selections by intersection with current selection.
        Ranks determined by the hypergeometric distribution.
        """
        dim_name = dataset.get_dim_name(0)
        sel_store = self._sel_stores[dim_name]
        selection_obj = self.project.get_selection()
        current_selection = selection_obj.get(dim_name)
        if current_selection==None: return

        pvals = hypergeom.gene_hypergeo_test(current_selection, dataset)

        for row in sel_store:
            if row[1]==dataset:
                for child in row.iterchildren():
                    name = child[0]
                    child[3] = pvals[name][0]
                    child[4] = pvals[name][1]
                    child[5] = pvals[name][2]
        
        sel_store.set_sort_column_id(5, gtk.SORT_ASCENDING)

    #
    # GTK callbacks
    #
    
    def _drag_data_received(self, widget, drag_context, x, y,
                              selection, info, timestamp):

        treestore, path = selection.tree_get_row_drag_data()
        i = treestore.get_iter(path)
        obj = treestore.get_value(i, 2)
        if isinstance(obj, dataset.CategoryDataset):
            self.add_dataset(obj)
            self._dimlist_controller.set_dimension(obj.get_dim_name(0))
        widget.emit_stop_by_name('drag-data-received')

    def _on_cursor_changed(self, widget):
        "Show the list of identifier strings."

        store = self._sel_stores[self._dimension]

        p = self._seltree.get_cursor()[0]
        i = store.get_iter(p)
        obj = store.get_value(i, 1)

        if isinstance(obj, dataset.Selection):
            self._idlist_controller.set_selection(obj)
        
    def _on_row_activated(self, widget, path, column):
        store = self._sel_stores[self._dimension]

        i = store.get_iter(path)
        obj = store.get_value(i, 1)
        if isinstance(obj, dataset.Dataset):
            seltree = self._seltree
            if seltree.row_expanded(path):
                seltree.collapse_row(path)
            else:
                seltree.expand_row(path, True)
        elif isinstance(obj, dataset.Selection):
            self.project.set_selection(self._dimension,
                                       obj[self._dimension])

    def _on_button_pressed(self, widget, event):
        """Button press callbak."""
        if event.button == 3:
            self._seltree_menu.popup(None, None, None, event.button, event.time)

    def _on_seltree_sort(self, menuitem):
        """Sort selection tree if row is category dataset."""
        store = self._sel_stores[self._dimension]
        p = self._seltree.get_cursor()[0]
        i = store.get_iter(p)
        obj = store.get_value(i, 1)
        if isinstance(obj, dataset.CategoryDataset):
            self._sort_selections(obj)

class DimListController:
    def __init__(self, dimlist, seltree_controller):
        
        self._current_dim = None
        self._seltree_controller = seltree_controller
        
        ## dimstore is a list of all dimensions in the application
        self.dimstore = gtk.ListStore(gobject.TYPE_STRING)

        ## The widgets we are controlling
        self.dimlist = dimlist

        ## Set up dimensions list
        dimlist.set_model(self.dimstore)
        
        renderer = gtk.CellRendererText()
        dim_column = gtk.TreeViewColumn('Dimension', renderer, text=0)
        dimlist.insert_column(dim_column, 0)

        dimlist.connect('row-activated', self._dim_row_activated)
        dimlist.connect('cursor-changed', self._dim_cursor_changed)

    ## 
    ## Public interface
    ##
    def set_project(self, project):
        """Dependency injection."""
        self.project = project
        self.dim_names = project.dim_names
        self.update_dims()
        project.add_dataset_observer(self)

    def get_dimension(self, dim):
        """Returns the iterator to the dimension with the given name, or
        None if not found."""
        
        i = self.dimstore.get_iter_first()
        while i:
            if self.dimstore.get_value(i, 0) == dim:
                return i
            i = self.dimstore.iter_next(i)
        return None

    def set_dimension(self, dimname):
        self._current_dim = dimname

        dim = self.get_dimension(self._current_dim)
        path = self.dimstore.get_path(dim)
        
        if self.dimlist.get_cursor()[0] != path:
            self.dimlist.set_cursor(self.dimstore.get_path(dim))

        self._seltree_controller.set_dimension(dimname)

    def dataset_changed(self):
        """Callback function from Project."""
        self.update_dims()

    ##
    ## Private interface
    ##
    def update_dims(self):
        """Update the list of dimensions shown"""
        for dim in self.dim_names:
            if not self.get_dimension(dim):
                self.dimstore.insert_after(None, (dim,))

    def _dimension_filter(self, store, row):
        """Filters out everything but the selected dimension."""
        row_dim = store.get_value(row, 2)
        return row_dim == self._current_dim

    ##
    ## GTK Callbacks.
    ##
    def _dim_cursor_changed(self, widget):
        cursor = self.dimlist.get_cursor()[0]
        i = self.dimstore.get_iter(cursor)
        row = self.dimstore.get_value(i, 0)
        self.set_dimension(row)

    def _dim_row_activated(self, widget, path, column):
        self._seltree_controller.set_dimension(dim)