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

import logger, dataset, main
import annotations
from lib import hypergeom


class SimpleMenu(gtk.Menu):
    def __init__(self):
        gtk.Menu.__init__(self)
        
    def add_simple_item(self, title, function, *args):
        item = gtk.MenuItem(title)
        item.connect('activate', function, *args)
        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()

        # Return if no selection
        if selection == None:
            return

        # Otherwise show selection, possibly with annotations.
        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:
            print "set_sensitive(False)"
            self._menu_ann.set_sensitive(False)
        else:
            annotations_menu = gtk.Menu()
            print "set_sensitive(True)"
            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)

    def import_annotation_file(self):
        """Pops up a file dialog and ask the user to select the annotation
        file to be loaded. Only one file can be selected. The file is loaded
        into a annotations.AnnotationDictHandler object"""

        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()

    ##
    ## 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):
        self.import_annotation_file()
        

class SelectionListController:
    def __init__(self, seltree, idlist_controller):
        self._seltree = seltree
        self._sel_stores = {}
        self._detail_cols = []
        self._dimension = None
        self._idlist_controller = idlist_controller
        self._details_on = False
        
        # Selection column
        renderer = gtk.CellRendererText()
        sel_column = gtk.TreeViewColumn('Selection', renderer, text=0)
        sel_column.set_resizable(True)
        sel_column.set_max_width(200)
        seltree.insert_column(sel_column, 0)

        # Detail columns
        cols = [('In CS', 3), ('All', 4), ('Rank', 5)]
        for name, store_col_num in cols:
            col = gtk.TreeViewColumn(name, renderer, text=store_col_num)
            col.set_sort_indicator(True)
            col.set_sort_column_id(store_col_num)
            col.set_sort_order(gtk.SORT_ASCENDING)
            
            self._detail_cols.append(col)
        # 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('Show details',
                                           self._enable_details, True)
        self._seltree_menu.add_simple_item('Hide details', 
                                           self._enable_details, False)

    #
    # Public interface
    #
    def activate(self):
        self._seltree.set_cursor((0,))

    def set_project(self, project):
        """Dependency injection."""
        main.project.add_selection_observer(self)
        
    def set_dimlist_controller(self, dimlist_controller):
        """Dependency injection of the dimension list controller."""
        self._dimlist_controller = dimlist_controller

    def set_dimension(self, dim):
        """Set the current dimension, changing the model of the treeview
        to match dim. After this the current dimension of the identifier list
        is updated."""
        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, 0)
                store.insert_after(None, None, values)
            else:
                # update size of current selection
                for row in store:
                    if row[1]==selection:
                        row[3] = row[4] = len(selection[dim])

        path = self._seltree.get_cursor()
        if path and self._sel_stores.has_key(self._dimension):
            it = self._sel_stores[self._dimension].get_iter(path[0])
            sel = self._sel_stores[self._dimension].get_value(it, 1)
            self._idlist_controller.set_selection(sel)

    def add_dataset(self, dataset):
        """Converts a CategoryDataset to Selection objects and adds it to
        the selection tree. The name of the dataset will be the parent 
        node in the tree, and the identifers along the first axis will
        be added as the names of the subselections."""
        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 = main.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 = main.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 _enable_details(self, widget, bool):
        if self._details_on == bool : return
        self._details_on = bool
        if bool==True:
            for col in self._detail_cols:
                self._seltree.insert_column(col, -1)
        else:
            for col in self._detail_cols:
                self._seltree.remove_column(col)
    
    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)
        else:
            self._idlist_controller.set_selection(None)
        
    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):
            main.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
        
        self.show_hidden = False

        ## dimstore is a list of all dimensions in the application
        self.dimstore = gtk.ListStore(gobject.TYPE_STRING)

        # filter for hiding dims prefixed with underscore
        self.dimstore_filter = self.dimstore.filter_new()
        self.dimstore_filter.set_visible_func(self._dimension_filter)
        
        ## The widgets we are controlling
        self.dimlist = dimlist

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

        # Signals
        dimlist.connect('row-activated', self._dim_row_activated)
        dimlist.connect('cursor-changed', self._dim_cursor_changed)
        dimlist.connect('button-press-event', self._dimlist_button_pressed)
        
        # Set up  dimension context menu
        self._dimlist_menu = SimpleMenu()
        self._dimlist_menu.add_simple_item('Hide', self._on_dim_hide)
        self._dimlist_menu.add_simple_item('Show all', self._on_dim_show)


    ## 
    ## 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_filter.get_iter_first()
        while i:
            if self.dimstore_filter.get_value(i, 0) == dim:
                return i
            i = self.dimstore_filter.iter_next(i)
        return None

    def set_dimension(self, dimname):
        """Sets the current dimension."""
        self._current_dim = dimname
        
        dim = self.get_dimension(self._current_dim)
        path = self.dimstore_filter.get_path(dim)
        
        if self.dimlist.get_cursor()[0] != path:
            self.dimlist.set_cursor(self.dimstore_filter.get_path(dim))
        self._seltree_controller.set_dimension(dimname)

    def dataset_changed(self):
        """Callback function from Project."""
        self.update_dims()
    
    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,))
        self.dimstore_filter.refilter()

    #
    # Private interface
    #
    def _dimension_filter(self, store, row):
        """Filters out dimensions with underscore prefix."""
        if self.show_hidden:
            return True

        visible = False
        name = store.get_value(row, 0)
        if name != None:
            visible = name[0]!="_"
        return visible

    #
    # GTK Callbacks.
    #
    def _on_dim_hide(self, menuitem):
        """Menu item callback function which hides underscore prefixed
        dimensions."""
        self.show_hidden = False
        self.dimstore_filter.refilter()

    def _on_dim_show(self, menuitem):
        """Menu item callback function that shows underscore prefixed
        dimension names."""
        self.show_hidden = True
        self.dimstore_filter.refilter()
    
    def _dim_cursor_changed(self, widget):
        cursor = self.dimlist.get_cursor()[0]
        i = self.dimstore_filter.get_iter(cursor)
        row = self.dimstore_filter.get_value(i, 0)
        self.set_dimension(row)
        self._seltree_controller.activate()

    def _dim_row_activated(self, widget, path, column):
        self._seltree_controller.set_dimension(dim)
        
    def _dimlist_button_pressed(self, widget, event):
        if event.button == 3:
            self._dimlist_menu.popup(None, None, None, event.button, event.time)