import sys
import cProfile
import pstats

import PyQt4.QtGui as QtGui
import PyQt4.QtCore as QtCore


#########################################################################
# Underlying data
# ----------------
# - RuntimeItems hold the data. They come from a database.
# - ViewItems are the objects, that are given to the model indexes of Qt. 
#   They are constructed according to some rules like filters and 
#   configuration.
# - DummieViewItemFactory processes the rules and configurations. 
#   The example here is simplfied. An instance of the factory is given 
#   to each ViewItem. 
#   The view item calls the 
#   DummieViewItemFactory.get_view_item_children method
#   to request calculation of its children on demand.
# - For this demo-version, the number of items is controlled by 
#   DummieViewItemFactory.max_items. It's passed in by the constructor. 
# - Nesting as high as possible: One child per parent.
#########################################################################


class RuntimeItem(object):
    """Represent the real world business items. These objects 
    have a lot of relations.
    """
    
    def __init__(self, name, ident, item_type):
        self.name = name
        self.ident = ident
        self.item_type = item_type

        
class ViewItem(object):
    """Represent items that are to be shown to the user in a QTreeView.
    Those items do only occur one time in a view. They have a 
    corresponding runtime_item.
    The children are calculated by the view_item_factory on demand.
    """
    
    def __init__(self, view_item_factory, runtime_item=None, parent=None, 
                 hidden_runtime_items=None):
        self.view_item_factory = view_item_factory
        self.runtime_item = runtime_item
        self.parent = parent
        self.hidden_runtime_items = hidden_runtime_items
        
    @property
    def children(self):
        try:
            return self._children
        except AttributeError:
            self._children = \
                self.view_item_factory.get_view_item_children(self)
            return self._children
        
    @children.setter
    def children(self, children):
        self._children = children

        
class DummieViewItemFactory(object):
    """Creates the view_items. This is a dumb dummie as a simple 
    example. Normally a lot of things happen here like filtering 
    and configuration. But once the view_item hierachy is build, 
    this shouldn't be called at all.
    """
    
    def __init__(self, runtime_item, max_items):            
        self.runtime_item = runtime_item
        self.max_items = max_items
        self.item_counter = 0
        self.aux_root_view_item = ViewItem(self)
            
    def get_view_item_children(self, view_item_parent):
        if self.item_counter > self.max_items:
            return []
        view_items = []
        for i in range(1):
            self.item_counter += 1
            view_items.append(ViewItem(self, self.runtime_item, view_item_parent))
        return view_items
    
    
#########################################################################
# Qt classes
# ----------------
# - This should be standard stuff. I've got most of it from the Rapid 
#   GUI Programming book.
# - The ActiveColums class tells the model which colums to use.
# - The TreeView has a context menu with navigation actions.
# - The expand_all calls the Qt slot. Here the surprise for the 
#   performance.
#########################################################################


class ActiveColumns(object):
    
    def __init__(self, columns):
        self.columns = columns


class TreeView(QtGui.QTreeView):

    def __init__(self, aux_root_view_item, active_columns, parent=None, 
                 header_hidden=False):
        super(TreeView, self).__init__(parent)
        self.setIndentation(10)
        self.active_columns = active_columns
        self.setAlternatingRowColors(True)
        self.setHeaderHidden(header_hidden)
        self.setAllColumnsShowFocus(True)
        self.setUniformRowHeights(True)
        self.setContextMenuPolicy(QtCore.Qt.ActionsContextMenu)
        
        self.model = TreeModel(aux_root_view_item, self)
        self.setModel(self.model)
        
        e_a_action = QtGui.QAction("Expand all", self)
        e_a_action.setToolTip("Expands all items of the tree.")
        e_a_action.triggered.connect(self.expand_all)
        
        e_a_b_action = QtGui.QAction("Expand all below", self)
        e_a_b_action.setToolTip("Expands all items under the selection.")
        e_a_b_action.triggered.connect(self.expand_all_below)
        
        c_a_action = QtGui.QAction("Collapse all", self)
        c_a_action.setToolTip("Collapses all items of the tree.")
        c_a_action.triggered.connect(self.collapse_all)
        
        c_a_b_action = QtGui.QAction("Collapse all below", self)
        c_a_b_action.setToolTip("Collapses all items under the selection.")
        c_a_b_action.triggered.connect(self.collapse_all_below)
        
        for action in (e_a_action, c_a_action, e_a_b_action, c_a_b_action):
            self.addAction(action)
            
    def expand_all(self):
        self.expandAll()
            
    def collapse_all(self):
        self.collapseAll()
            
    def expand_all_below(self):
        def expand_all_below_recursive(parent_index):
            self.expand(parent_index)
            children_indexes = \
                self.model.get_children_indexes(parent_index)
            for child_index in children_indexes:
                expand_all_below_recursive(child_index)
 
        indexes = self.selectedIndexes()
        if indexes:
            index = indexes[0]
            expand_all_below_recursive(index)
                
    def collapse_all_below(self):
        def collapse_all_below_recursive(parent_index):
            self.collapse(parent_index)
            children_indexes = \
                self.model.get_children_indexes(parent_index)
            for child_index in children_indexes:
                collapse_all_below_recursive(child_index)            

        indexes = self.selectedIndexes()
        if indexes:
            index = indexes[0]
            collapse_all_below_recursive(index)
        
class TreeModel(QtCore.QAbstractItemModel):

    def __init__(self, aux_root_view_item, parent):
        super(TreeModel, self).__init__(parent)
        self.aux_root_view_item = aux_root_view_item
        self.active_columns = parent.active_columns

    def rowCount(self, parent_index):
        parent_view_item = self.view_item_from_index(parent_index)
        if parent_view_item is None:
            return 0
        return len(parent_view_item.children)
    
    def get_children_indexes(self, parent_index):
        children_indexes = []
        for row_no in range(self.rowCount(parent_index)):
            children_indexes.append(self.index(row_no, 0, parent_index))
        return children_indexes

    def columnCount(self, parent):
        return len(self.active_columns.columns)

    def data(self, index, role):
        if role == QtCore.Qt.TextAlignmentRole:
            return int(QtCore.Qt.AlignTop|QtCore.Qt.AlignLeft)
        if role != QtCore.Qt.DisplayRole:
            return None
        view_item = self.view_item_from_index(index)
        try:
            data = getattr(view_item.runtime_item, 
                           self.active_columns.columns[index.column()])
        except AttributeError:
            data = ""
        return data

    def headerData(self, section, orientation, role):
        if (orientation == QtCore.Qt.Horizontal and
            role == QtCore.Qt.DisplayRole):
            assert 0 <= section <= len(self.active_columns.columns)
            return self.active_columns.columns[section]
        return QtCore.QVariant()

    def index(self, row, column, parent_index):
        view_item_parent = self.view_item_from_index(parent_index)
        return self.createIndex(row, column, 
                                view_item_parent.children[row])

    def parent(self, child_index):
        child_view_item = self.view_item_from_index(child_index)
        if child_view_item is None:
            return QtCore.QModelIndex()
        parent_view_item = child_view_item.parent
        if parent_view_item is None:
            return QtCore.QModelIndex()
        grandparent_view_item = parent_view_item.parent
        if grandparent_view_item is None:
            return QtCore.QModelIndex()
        grandparent_view_item
        row = grandparent_view_item.children.index(parent_view_item)
        assert row != -1
        return self.createIndex(row, 0, parent_view_item)
    
    def view_item_from_index(self, index):
        return (index.internalPointer()
                if index.isValid() else self.aux_root_view_item)
    
    def canFetchMore(self,parent_index):
        parent_item = self.view_item_from_index(parent_index)
        if parent_item.children:
            return True
        return False
    
    def fetchMore(self, parent_index):
        parent_item = self.view_item_from_index(parent_index)
        parent_item.children
    
    
if __name__ == "__main__":
    
    sys.setrecursionlimit(1000000)
    
    run_time_item = RuntimeItem("Test", "test_12", "Test Item")
    view_factory = DummieViewItemFactory(run_time_item, max_items=5000)
    active_colums = ActiveColumns(["name", "ident", "item_type"])

    app = QtGui.QApplication(sys.argv)
    tree_view = TreeView(view_factory.aux_root_view_item, active_colums)
    app.setApplicationName("IPDM")
    tree_view.show()
    #app.exec_()
    cProfile.run('app.exec_()', 'profdata')
    p = pstats.Stats('profdata')
    p.sort_stats('time').print_stats()
