from PyQt5 import QtCore, QtWidgets, QtGui
from PyQt5 import uic

from itertools import groupby
import sys



# ------------------------------------------------------------------------------
# | CategoryNode                                                               |
# |                                                                            |
# |                                                                            |
# |                                                                            |
# | Defines the nodes of the data tree that underlies CategoryTreeModel        |
# ------------------------------------------------------------------------------
class CategoryNode:
    def __init__(self, desc, parent=None):
        self.parent_node = parent
        self._desc = desc
        self.children = []

    # ----------------------------------------------------------------------
    # | CategoryNode accessor functions
    # ----------------------------------------------------------------------

    # Number of columns is always one in this model
    @staticmethod
    def n_columns():
        return 1

    # Return category description
    def desc(self):
        return self._desc

    # Set category description
    def set_desc(self, desc):
        self._desc = desc

    # Return row 
    def row(self):
        if self.parent_node:
            return self.parent_node.children.index(self)
        else:
            return 0

    # Return parent
    def parent(self):
        return self.parent_node

    # Add a child
    def append_child(self, c):
        self.children.append(c)

    # Return child from given row
    def child(self, row):
        return self.children[row]

    # Return number of children
    def n_children(self):
        return len(self.children)

    # Delete child from given row
    def del_child(self, row):
        del self.children[row]

    # ----

    # Printable representation
    def __repr__(self):
        return str(f'->{self._desc}')

    # Recursively dump tree contents
    def dump(self, i=0):
        print(f'''{(' ' * (i * 4))} {self.row()} -> {self}''')
        for c in self.children:
            c.dump(i+1)


# ------------------------------------------------------------------------------
# | CategoryTreeModel                                                          |
# |                                                                            |
# |                                                                            |
# |                                                                            |
# | Model used by CategoryDialog tree view                                     |
# ------------------------------------------------------------------------------
class CategoryTreeModel(QtCore.QAbstractItemModel):

    # ------------------------------------------------------------[override]
    # | __init__()
    # ----------------------------------------------------------------------
    def __init__(self, category_data, parent=None):
        super().__init__(parent)

        # Create root node
        self.root = CategoryNode('Category')

        # Create main category/subcategory node structure
        for g in groupby(category_data, lambda rec: rec[1]):
            sub_list = list(g[1])
            main_node = CategoryNode(g[0], self.root)
            self.root.append_child(main_node)
            for s in sub_list:
                if s[2] is not None:
                    sub_node = CategoryNode(s[2], main_node)
                    main_node.append_child(sub_node)

    # ------------------------------------------------------------[override]
    # | rowCount()
    # ----------------------------------------------------------------------
    def rowCount(self, parent=QtCore.QModelIndex()):
        if parent.isValid():
            return parent.internalPointer().n_children()
        else:
            return self.root.n_children()

    # ------------------------------------------------------------[override]
    # | columnCount()
    # ----------------------------------------------------------------------
    def columnCount(self, parent=QtCore.QModelIndex()):
        if parent.isValid():
            return parent.internalPointer().n_columns()
        else:
            return self.root.n_columns()

    # ------------------------------------------------------------[override]
    # | data()
    # ----------------------------------------------------------------------
    def data(self, index, role=QtCore.Qt.DisplayRole):
        if index.isValid():
            if role == QtCore.Qt.DisplayRole:
                return index.internalPointer().desc()
            elif role == QtCore.Qt.FontRole:
                font = QtGui.QFont('Candara')
                font.setPointSize(11)
                return font
            elif role == QtCore.Qt.ForegroundRole:
                color = '#1e62ce'
                return QtGui.QBrush(QtGui.QColor(color))

        return None

    # ------------------------------------------------------------[override]
    # | flags()
    # ----------------------------------------------------------------------
    def flags(self, index):
        if index.isValid():
            return QtCore.Qt.ItemIsEnabled | QtCore.Qt.ItemIsEditable
        else:
            return QtCore.Qt.NoItemFlags

    # ------------------------------------------------------------[override]
    # | index()
    # ----------------------------------------------------------------------
    def index(self, row, column, parent=QtCore.QModelIndex()):

        # According to legend, index() may sometimes be called with invalid parameters
        # (that is, with parameters that don't actually have an index).  So this is
        # necessary to keep it from crashing.
        # https://stackoverflow.com/questions/26680168/pyqt-treeview-index-error-removing-last-row
        if not self.hasIndex(row, column, parent):
            return QtCore.QModelIndex()

        elif parent.isValid():
            return self.createIndex(row, column, parent.internalPointer().child(row))
        else:
            return self.createIndex(row, column, self.root.child(row))

    # ------------------------------------------------------------[override]
    # | parent()
    # ----------------------------------------------------------------------
    def parent(self, index):
        if not index.isValid():
            return QtCore.QModelIndex()

        child_node = index.internalPointer()
        parent_node = child_node.parent()
        if parent_node == self.root:
            r = QtCore.QModelIndex()
        else:
            r = self.createIndex(parent_node.row(), 0, parent_node)
        return r


    # ------------------------------------------------------------[override]
    # | setData()
    # ----------------------------------------------------------------------
    def setData(self, index, value, role=QtCore.Qt.EditRole):
        if role == QtCore.Qt.EditRole:

            # Set value in data tree node
            node = index.internalPointer()
            node.set_desc(value)

        return True


    # ------------------------------------------------------------[override]
    # | insertRows()
    # ----------------------------------------------------------------------
    def insertRows(self, row, rows, parent=QtCore.QModelIndex()):
        assert rows == 1

        # Determine parent
        if parent.internalPointer() is None:
            parent_node = self.root
        else:
            parent_node = parent.internalPointer()

        # Create new data tree node and add it to parent
        new = CategoryNode(None, parent_node)
        self.beginInsertRows(parent, row, row + rows - 1)
        parent_node.append_child(new)
        self.endInsertRows()

        return True

    # ------------------------------------------------------------[override]
    # | removeRows()
    # ----------------------------------------------------------------------
    def removeRows(self, row, rows, parent=QtCore.QModelIndex()):
        assert rows == 1

        # Determine parent
        if parent.internalPointer() is None:
            parent_node = self.root
        else:
            parent_node = parent.internalPointer()

        # Delete it from the parent
        self.beginRemoveRows(parent, row, row + rows - 1)
        parent_node.del_child(row)
        self.endRemoveRows()

        return True


# ------------------------------------------------------------------------------
# | Proxy                                                                      |
# |                                                                            |
# |                                                                            |
# |                                                                            |
# | QSortFilterProxyModel object                                               |
# ------------------------------------------------------------------------------
class Proxy(QtCore.QSortFilterProxyModel):

    # ------------------------------------------------------------[override]
    # | __init__()
    # ----------------------------------------------------------------------
    def __init__(self, parent=None):
        super().__init__(parent)

    # ------------------------------------------------------------[override]
    # | setData()
    # ----------------------------------------------------------------------
    def setData(self, index, value, role=QtCore.Qt.EditRole):

        # Call source model's setData()
        super().setData(index, value, role)

        # Re-sort
        self.sort(0, QtCore.Qt.AscendingOrder)
        self.invalidate()


# ------------------------------------------------------------------------------
# | TreeView                                                                   |
# |                                                                            |
# |                                                                            |
# |                                                                            |
# | QTreeView object                                                           |
# ------------------------------------------------------------------------------
class TreeView(QtWidgets.QTreeView):
    def __init__(self, parent=None):
        super().__init__(parent)

    def closeEditor(self, editor, hint):
        proxy_model = self.model()
        proxy_idx = self.currentIndex()
        if not proxy_idx.data():
            source_model = proxy_model.sourceModel()
            source_idx = proxy_model.mapToSource(proxy_idx)
            source_model.removeRows(source_idx.row(), 1, source_model.parent(source_idx))

        super().closeEditor(editor, hint)


# ------------------------------------------------------------------------------
# | Dlg                                                                        |
# |                                                                            |
# |                                                                            |
# |                                                                            |
# | Dialog widget (defined in dlg.ui)                                          |
# ------------------------------------------------------------------------------
dlg_base, dlg_form = uic.loadUiType('dlg.ui')
class Dlg(dlg_base, dlg_form):

    def __init__(self, parent=None):
        super(dlg_base, self).__init__(parent)
        self.setupUi(self)
        self.setWindowFlags(QtCore.Qt.Dialog|QtCore.Qt.FramelessWindowHint)
        self.treeView.setHeaderHidden(True)
        self.treeView.setAnimated(True)

        # Set model and proxy
        category_data = [
            (4, 'Auto', None), (10, 'Auto', 'Gas'), (9, 'Auto', 'License'),
                (11, 'Auto', 'Service'),
            (23, 'Beer', None),
            (12, 'Clothing', None), (13, 'Clothing', 'Cleaning'),
            (18, 'Dining', None), (19, 'Dining', 'Breakfast'), (21, 'Dining', 'Lunch'),
            (14, 'Healthcare', None), (26, 'Healthcare', 'Dental'),
            (3, 'Household', None),
            (5, 'Pet Expenses', None), (7, 'Pet Expenses', 'Biskies'),
                (17, 'Pet Expenses', 'Chow'), (8, 'Pet Expenses', 'Grooming'),
                (6, 'Pet Expenses', 'Medical'), (15, 'Pet Expenses', 'Miscellaneous'),
            (25, 'Reimbursible', None)
        ]
        self.model = CategoryTreeModel(category_data)
        self.proxy = Proxy()
        self.proxy.setSourceModel(self.model)
        self.treeView.setModel(self.proxy)

        # Set Delegate
        self.delegate = CategoryDelegate(self.treeView)
        self.treeView.setItemDelegate(self.delegate)

        # Connect buttons
        self.addMainButton.clicked.connect(self.add_main)
        self.addSubButton.clicked.connect(self.add_sub)
        self.deleteButton.clicked.connect(self.delete)
        self.doneButton.clicked.connect(quit)

    def keyPressEvent(self, event):
        key = event.key()
        mod = int(event.modifiers())
        if key == QtCore.Qt.Key_Q and mod == QtCore.Qt.CTRL:
            exit()

    def add_main(self):
        print('Add Main')
        view = self.treeView
        model = view.model()
        model.insertRows(model.rowCount(), 1)
        index = model.index(model.rowCount() - 1, 0)
        view.setCurrentIndex(index)
        view.scrollTo(index, QtWidgets.QAbstractItemView.EnsureVisible)
        view.edit(index)

    def add_sub(self):
        print('Add Sub')
        view = self.treeView
        model = view.model()

        current = view.currentIndex()
        parent = model.parent(current)
        add_to = parent if parent.isValid() else current

        model.insertRows(model.rowCount(add_to), 1, add_to)
        index = model.index(model.rowCount(add_to) - 1, 0, add_to)
        view.setCurrentIndex(index)
        view.scrollTo(index, QtWidgets.QAbstractItemView.EnsureVisible)
        view.edit(index)

    def delete(self):
        print('Delete')
        view = self.treeView
        model = view.model()
        current = view.currentIndex()
        model.removeRows(current.row(), 1, model.parent(current))


# ------------------------------------------------------------------------------
# | CategoryDelegate                                                           |
# |                                                                            |
# |                                                                            |
# |                                                                            |
# | Delegate for editing                                                       |
# ------------------------------------------------------------------------------
class CategoryDelegate(QtWidgets.QStyledItemDelegate):

    def __init__(self, parent=None):
        super().__init__(parent)

    def createEditor(self, parent, option, index):
        self.editor = QtWidgets.QLineEdit(parent)
        return self.editor

    def setEditorData(self, editor, index):
        pass

    def updateEditorGeometry(self, editor, option, index):
        editor.setGeometry(option.rect)

    def setModelData(self, editor, model, index):
        model.setData(index, editor.text())


# ------------------------------------------------------------------------------
# | main()                                                                     |
# |                                                                            |
# |                                                                            |
# |                                                                            |
# |                                                                            |
# ------------------------------------------------------------------------------
if __name__ == '__main__':
    app = QtWidgets.QApplication(sys.argv)
    app.setStyle('Fusion')
    dlg = Dlg()
    dlg.show()
    sys.exit(app.exec_())
