nmreval/src/gui_qt/data/datawidget/datawidget.py

from __future__ import annotations

from nmreval.lib.colors import available_cycles

from .properties import PropWidget
from ...Qt import QtWidgets, QtGui, QtCore
from ..._py.datawidget import Ui_DataWidget
from ...lib.iconloading import make_action_icons
from ...lib.delegates import HeaderDelegate


class DataTree(QtWidgets.QTreeWidget):
    stateChanged = QtCore.pyqtSignal(list, list)  # selected, deselected
    keyChanged = QtCore.pyqtSignal(str, str)  # id, text
    positionChanged = QtCore.pyqtSignal(QtWidgets.QTreeWidgetItem)
    deleteItem = QtCore.pyqtSignal(list)
    moveItem = QtCore.pyqtSignal(list, str, str, int)  # items, from, to, new row
    copyItem = QtCore.pyqtSignal(list, str)
    saveFits = QtCore.pyqtSignal(list)
    extendFits = QtCore.pyqtSignal(list, bool)

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

        self.invisibleRootItem().setFlags(self.invisibleRootItem().flags() ^ QtCore.Qt.ItemIsDropEnabled)

        self.itemChanged.connect(self.data_change)
        self.itemClicked.connect(self.new_selection)

        self.setColumnCount(2)

        self.setSizePolicy(QtWidgets.QSizePolicy.MinimumExpanding, QtWidgets.QSizePolicy.Expanding)
        self.setDragDropMode(QtWidgets.QTreeView.InternalMove)
        self.setDefaultDropAction(QtCore.Qt.IgnoreAction)
        self.setSelectionMode(QtWidgets.QTreeView.ExtendedSelection)
        self.setSelectionBehavior(QtWidgets.QTreeView.SelectRows)

        self._checked_graphs = set()
        self._checked_sets = set()
        self.management = None

        header = QtWidgets.QHeaderView(QtCore.Qt.Horizontal, self)
        self.setHeader(header)
        header.setSectionResizeMode(0, QtWidgets.QHeaderView.Stretch)
        header.setVisible(False)
        header.moveSection(1, 0)
        self.setItemDelegateForColumn(1, HeaderDelegate())

    def add_graph(self, idd: str, name: str):
        item = QtWidgets.QTreeWidgetItem()
        item.setFlags(
            QtCore.Qt.ItemFlag.ItemIsSelectable |
            QtCore.Qt.ItemFlag.ItemIsDropEnabled |
            QtCore.Qt.ItemFlag.ItemIsEditable |
            QtCore.Qt.ItemFlag.ItemIsEnabled |
            QtCore.Qt.ItemFlag.ItemIsUserCheckable
        )
        item.setText(0, name)
        item.setData(0, QtCore.Qt.ItemDataRole.UserRole, idd)
        item.setCheckState(0, QtCore.Qt.CheckState.Checked)

        self.addTopLevelItem(item)
        self._checked_graphs.add(idd)
        item.setExpanded(True)

        self.update_indexes()

    def add_item(self, items: (tuple | list[tuple]), gid: str, update: bool = True):
        if isinstance(items, tuple):
            items = [items]

        for row in range(self.invisibleRootItem().childCount()):
            graph = self.invisibleRootItem().child(row)
            if graph.data(0, QtCore.Qt.ItemDataRole.UserRole) == gid:
                for (idd, name, value) in items:
                    item = QtWidgets.QTreeWidgetItem([name])
                    item.setToolTip(0, f'Value: {value}')
                    item.setData(0, QtCore.Qt.ItemDataRole.UserRole, idd)
                    item.setCheckState(0, QtCore.Qt.CheckState.Checked)
                    item.setFlags(
                        QtCore.Qt.ItemFlag.ItemIsSelectable |
                        QtCore.Qt.ItemFlag.ItemIsDragEnabled |
                        QtCore.Qt.ItemFlag.ItemIsEditable |
                        QtCore.Qt.ItemFlag.ItemIsEnabled |
                        QtCore.Qt.ItemFlag.ItemIsUserCheckable
                    )
                    graph.addChild(item)
                    self._checked_sets.add(idd)

                self.resizeColumnToContents(0)
                break
        if update:
            self.update_indexes()

    @QtCore.pyqtSlot(QtWidgets.QTreeWidgetItem)
    def data_change(self, item: QtWidgets.QTreeWidgetItem, emit: bool = True) -> tuple[set, set]:
        idd = item.data(0, QtCore.Qt.ItemDataRole.UserRole)
        is_selected = item.checkState(0) == QtCore.Qt.CheckState.Checked
        to_be_hidden = set()
        to_be_shown = set()

        # item is top-level item / graph
        if item.parent() is None:
            was_selected = idd in self._checked_graphs

            # check state changed to selected
            if is_selected != was_selected:
                # check state changed to checked
                if is_selected:
                    self._checked_graphs.add(idd)
                    iterator = QtWidgets.QTreeWidgetItemIterator(item)
                    iterator += 1
                    self.blockSignals(True)
                    for i in range(item.childCount()):
                        child = item.child(i)
                        child.setCheckState(0, QtCore.Qt.CheckState.Checked)
                        to_be_shown.add(child.data(0, QtCore.Qt.ItemDataRole.UserRole))
                        self._checked_sets.add(child.data(0, QtCore.Qt.ItemDataRole.UserRole))
                    self.blockSignals(False)

                # check state change to unchecked
                else:
                    self._checked_graphs.remove(idd)
                    self.blockSignals(True)
                    for i in range(item.childCount()):
                        child = item.child(i)
                        child.setCheckState(0, QtCore.Qt.CheckState.Unchecked)
                        to_be_hidden.add(child.data(0, QtCore.Qt.ItemDataRole.UserRole))
                        try:
                            self._checked_sets.remove(child.data(0, QtCore.Qt.ItemDataRole.UserRole))
                        except KeyError:
                            pass
                    self.blockSignals(False)

            else:
                self.keyChanged.emit(idd, item.text(0))

        # item is a set
        else:
            was_selected = idd in self._checked_sets
            if is_selected != was_selected:
                if is_selected:
                    to_be_shown.add(idd)
                    self._checked_sets.add(idd)

                else:
                    to_be_hidden.add(idd)
                    try:
                        self._checked_sets.remove(idd)
                    except KeyError:
                        pass

            else:
                if emit:
                    self.keyChanged.emit(idd, item.text(0))

        if (to_be_shown or to_be_hidden) and emit:
            self.stateChanged.emit(list(to_be_shown), list(to_be_hidden))

        return to_be_shown, to_be_hidden

    @QtCore.pyqtSlot(QtWidgets.QTreeWidgetItem)
    def new_selection(self, item: QtWidgets.QTreeWidgetItem):
        if item.parent() is None:
            self.management.select_window(item.data(0, QtCore.Qt.ItemDataRole.UserRole))

    def dropEvent(self, evt: QtGui.QDropEvent):
        dropped_index = self.indexAt(evt.pos())
        if not dropped_index.isValid():
            return

        to_parent = self.itemFromIndex(dropped_index).parent()
        append = False
        if not to_parent:
            # dropped on graph item -> append items
            to_parent = self.itemFromIndex(dropped_index)
            append = True

        # index may change
        persistent_drop = QtCore.QPersistentModelIndex(dropped_index)

        tobemoved = []
        take_from = []
        for it in self.selectedItems():
            from_parent = it.parent()
            if from_parent is None:
                continue

            from_parent.removeChild(it)
            tobemoved.append(it)
            take_from.append(from_parent.data(0, QtCore.Qt.ItemDataRole.UserRole))

        pos = QtCore.QModelIndex(persistent_drop)
        if self.dropIndicatorPosition() == QtWidgets.QAbstractItemView.BelowItem:
            pos = pos.sibling(pos.row()+1, 0)

        row = pos.row()
        if (row == -1) or append:
            to_parent.addChildren(tobemoved)
        else:
            to_parent.insertChildren(row, tobemoved)

        self.management.move_sets([it.data(0, QtCore.Qt.ItemDataRole.UserRole) for it in tobemoved],
                                  to_parent.data(0, QtCore.Qt.ItemDataRole.UserRole), take_from,
                                  pos=-1 if append else row)

        self.update_indexes()

    def move_sets(self, sid: str, gid_in: str, gid_out: str):
        self.blockSignals(True)
        to_parent = None
        from_parent = None
        it = None

        iterator = QtWidgets.QTreeWidgetItemIterator(self)
        while iterator.value():
            item = iterator.value()
            if item is not None:
                data = item.data(0, QtCore.Qt.ItemDataRole.UserRole)
                if data == gid_out:
                    from_parent = item

                elif data == gid_in:
                    to_parent = item

                elif data == sid:
                    it = item

                iterator += 1

        if (from_parent is None) or (to_parent is None) or (it is None):
            print('Komisch')
            return

        from_parent.removeChild(it)
        to_parent.addChild(it)

        self.update_indexes()

        self.blockSignals(False)

    def sort(self, graph_item: QtWidgets.QTreeWidgetItem, mode: str = 'value'):
        graph_id = graph_item.data(0, QtCore.Qt.ItemDataRole.UserRole)
        sets = self.management.get_attributes(graph_id, mode)
        sets = [el[0] for el in sorted(sets.items(), key=lambda x: x[1])]

        self.management.move_sets(sets, graph_id, graph_id, pos=-1)

        self.blockSignals(True)

        children = graph_item.takeChildren()

        for s in sets:
            for c in children:
                if c.data(0, QtCore.Qt.ItemDataRole.UserRole) == s:
                    graph_item.addChild(c)

        self.update_indexes()

        self.blockSignals(False)

    def update_indexes(self):
        graph_cnt = -1
        set_cnt = 0
        iterator = QtWidgets.QTreeWidgetItemIterator(self)
        self.blockSignals(True)
        while iterator.value():
            item = iterator.value()
            if item is not None:
                if item.parent() is None:
                    graph_cnt += 1
                    set_cnt = 0
                    item.setText(1, f'G[{graph_cnt}]')
                else:
                    item.setText(1, f'.S[{set_cnt}]')
                    set_cnt += 1

            iterator += 1

        self.resizeColumnToContents(1)
        self.blockSignals(False)

    def set_name(self, sid, name):
        iterator = QtWidgets.QTreeWidgetItemIterator(self)
        while iterator.value():
            item = iterator.value()
            if item is not None:
                data = item.data(0, QtCore.Qt.ItemDataRole.UserRole)
                if data == sid:
                    if name != item.text(0):
                        item.setText(0, name)
                    break

                iterator += 1

    def keyPressEvent(self, evt: QtGui.QKeyEvent):
        if evt.key() == QtCore.Qt.Key.Key_Delete:
            rm_sets = []
            rm_graphs = []
            for idx in self.selectedIndexes():
                if idx.column() == 1:
                    continue

                item = self.itemFromIndex(idx)
                if item.parent() is None:
                    for c_i in range(item.childCount()):
                        # add sets inside graph to removal
                        child_data = item.child(c_i).data(0, QtCore.Qt.ItemDataRole.UserRole)
                        if child_data not in rm_sets:
                            rm_sets.append(child_data)
                    rm_graphs.append(item.data(0, QtCore.Qt.ItemDataRole.UserRole))

                else:
                    item_data = item.data(0, QtCore.Qt.ItemDataRole.UserRole)
                    if item_data not in rm_sets:
                        rm_sets.append(item_data)

            # self.deleteItem.emit(rm_sets+rm_graphs)
            self.management.delete_sets(rm_sets+rm_graphs)

        elif evt.key() == QtCore.Qt.Key.Key_Space:
            sets = []
            from_parent = []

            for idx in self.selectedIndexes():
                if idx.column() != 0:
                    continue
                item = self.itemFromIndex(idx)
                if item.parent() is None:
                    for i in range(item.childCount()):
                        child = item.child(i)
                        from_parent.append(child)
                sets.append(item)

            to_be_hidden = set()
            to_be_shown = set()
            self.blockSignals(True)
            for it in sets:
                if it in from_parent:
                    continue
                it.setCheckState(0, QtCore.Qt.CheckState.Unchecked if it.checkState(0) == QtCore.Qt.CheckState.Checked else QtCore.Qt.CheckState.Checked)
                s1, s2 = self.data_change(it, emit=False)
                to_be_hidden |= s2
                to_be_shown |= s1
            self.blockSignals(False)
            self.stateChanged.emit(list(to_be_shown), list(to_be_hidden))
        else:
            super().keyPressEvent(evt)

    def mousePressEvent(self, evt: QtGui.QMouseEvent):
        # disable drag-and-drop when column 1 ('header') is clicked
        idx = self.indexAt(evt.pos())
        self.setDragEnabled(idx.column() == 0)
        super().mousePressEvent(evt)

    def remove_item(self, ids: list[str]):
        iterator = QtWidgets.QTreeWidgetItemIterator(self)

        toberemoved = []
        graph_removal = []

        # find all items that have to be removed
        while iterator.value():
            item = iterator.value()
            _id = item.data(0, QtCore.Qt.ItemDataRole.UserRole)
            if _id in ids:
                try:
                    item_parent = item.parent()
                    if item_parent is None:
                        raise AttributeError

                    idx = item_parent.indexOfChild(item)
                    # item.parent().takeChild(idx)

                    toberemoved.append((item_parent, idx))

                    if _id in self._checked_sets:
                        self._checked_sets.remove(_id)

                except AttributeError:
                    idx = self.invisibleRootItem().indexOfChild(item)
                    # self.invisibleRootItem().takeChild(idx)

                    graph_removal.append(idx)

                    if _id in self._checked_graphs:
                        self._checked_graphs.remove(_id)

            iterator += 1

        for (item, set_idx) in sorted(toberemoved, key=lambda x: x[1], reverse=True):
            item.takeChild(set_idx)

        for graph_idx in sorted(graph_removal, reverse=True):
            self.invisibleRootItem().takeChild(graph_idx)

        self.update_indexes()

    def contextMenuEvent(self, evt):
        menu = QtWidgets.QMenu()
        _ = menu.addAction('Hello')
        _.setEnabled(False)
        menu.addSeparator()

        if self.invisibleRootItem().childCount() == 0 and len(self.selectedIndexes()) == 0:
            rdn_action = menu.addAction('Randomness')

            action = menu.exec(evt.globalPos())

            if action == rdn_action:
                import webbrowser
                webbrowser.open('https://en.wikipedia.org/wiki/Special:Random')

        elif all([self.itemFromIndex(i).parent() is None for i in self.selectedIndexes()]):
            # only graphs selected
            self.ctx_graphs(evt, menu)

        else:
            self.ctx_sets(evt, menu)

    def ctx_graphs(self, evt, menu):
        copy_action = menu.addAction('Replicate graph!')
        del_action = menu.addAction('Exterminate graph!')

        sort_menu = menu.addMenu('Sort sets')
        for label in ['By name', 'By value']:
            sort_menu.addAction(label)

        col_menu = menu.addMenu('Color cycle')
        for c in available_cycles.keys():
            col_menu.addAction(c)

        action = menu.exec(evt.globalPos())
        if action is None:
            return

        graphs = []
        items = []
        for i in self.selectedIndexes():
            if i.column() == 0:
                continue
            items.append(self.itemFromIndex(i))
            graphs.append(self.itemFromIndex(i).data(0, QtCore.Qt.ItemDataRole.UserRole))

        if action == del_action:
            for gid in graphs:
                self.management.delete_graph(gid)

        elif action == copy_action:
            for gid in graphs:
                self.management.copy_graph(gid)

        elif action.parent() == col_menu:
            for gid in graphs:
                self.management.set_cycle(self.management.graphs[gid].sets, action.text())

        elif action.parent() == sort_menu:
            for i in items:
                self.sort(i, mode=action.text().split()[1])

        evt.accept()

    def ctx_sets(self, evt, menu):
        del_action = menu.addAction('Exterminate sets')
        cp_action = menu.addAction('Replicate sets')
        cat_action = menu.addAction('Join us!')
        plt_action = save_action = extend_action = subfit_action = None
        menu.addSeparator()
        col_menu = menu.addMenu('Color cycle')
        for c in available_cycles.keys():
            col_menu.addAction(c)

        idx = {}
        has_fits = False
        for i in self.selectedIndexes():
            if i.column() == 0:
                continue

            item = self.itemFromIndex(i)
            parent = item.parent()
            if parent is None:
                continue

            else:
                graph_id = parent.data(0, QtCore.Qt.ItemDataRole.UserRole)
                if graph_id not in idx:
                    idx[graph_id] = []
                # collect sets in their graph
                idx[graph_id].append(item.data(0, QtCore.Qt.ItemDataRole.UserRole))
                data = self.management[item.data(0, QtCore.Qt.ItemDataRole.UserRole)]
                if data.mode == 'fit':
                    has_fits = True

        if has_fits:
            menu.addSeparator()
            plt_action = menu.addAction('Plot fit parameter')
            save_action = menu.addAction('Save fit parameter')
            extend_action = menu.addAction('Extrapolate fit')
            subfit_action = menu.addAction('Plot partial functions')

        action = menu.exec(evt.globalPos())

        s = []
        for gid, sets in idx.items():
            s.extend(sets)

        if action is None:
            return

        if action == del_action:
            self.management.delete_sets(s)

        elif action == cp_action:
            for gid, sets in idx.items():
                self.management.copy_sets(sets, gid)

        elif action == cat_action:
            self.management.cat(s)

        elif action == plt_action:
            self.management.make_fit_parameter(s)

        elif action == save_action:
            self.saveFits.emit(s)

        elif action == extend_action:
            self.extendFits.emit(s, False)

        elif action == subfit_action:
            self.extendFits.emit(s, True)

        elif action.parent() == col_menu:
            self.management.set_cycle(s, action.text())

        evt.accept()

    def highlight(self, gid: str):
        iterator = QtWidgets.QTreeWidgetItemIterator(self)
        while iterator.value():
            item = iterator.value()
            if item is not None:
                if item.data(0, QtCore.Qt.ItemDataRole.UserRole) == gid:
                    item.setBackground(0, QtGui.QBrush(QtGui.QColor('gray')))
                else:
                    item.setBackground(0, QtGui.QBrush())
                iterator += 1

    def uncheck_sets(self, sets: list[str]):
        self.blockSignals(True)
        iterator = QtWidgets.QTreeWidgetItemIterator(self)
        self._checked_sets = set()
        while iterator.value():
            item = iterator.value()
            if item is not None:
                if item.data(0, QtCore.Qt.ItemDataRole.UserRole) in sets:
                    item.setCheckState(0, QtCore.Qt.CheckState.Unchecked)
                else:
                    self._checked_sets.add(item.data(0, QtCore.Qt.ItemDataRole.UserRole))
                iterator += 1
        self.blockSignals(False)


class DataWidget(QtWidgets.QWidget, Ui_DataWidget):
    keyChanged = QtCore.pyqtSignal(str, str)
    deleteItem = QtCore.pyqtSignal(list)
    startShowProperty = QtCore.pyqtSignal(list)
    propertyChanged = QtCore.pyqtSignal(list, str, str, object)

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

        self.setupUi(self)
        self.tree = DataTree(self)
        self.verticalLayout.addWidget(self.tree)
        # noinspection PyUnresolvedReferences
        self.tree.selectionModel().selectionChanged.connect(lambda x, y: self.show_property(x))

        self.tree.keyChanged.connect(lambda x, y: self.keyChanged.emit(x, y))

        self.proptable = PropWidget(self)
        self.propwidget.addWidget(self.proptable)
        self.propwidget.setText('Properties')
        self.propwidget.expansionChanged.connect(self.show_property)
        self.proptable.propertyChanged.connect(self.change_property)

        self.pokemon_toolbutton.clicked.connect(self.catchthemall)

        make_action_icons(self)

    def add_graph(self, idd: str, name: str):
        self.tree.blockSignals(True)
        self.tree.add_graph(idd, name)
        self.tree.blockSignals(False)

    def add_item(self, idd: str, name: str, value: str, gid: str, update: bool= True):
        self.tree.blockSignals(True)
        self.tree.add_item((idd, name, value), gid, update=update)
        self.tree.blockSignals(False)

    def add_item_list(self, loi: list, gid: str):
        self.tree.blockSignals(True)
        self.tree.add_item(loi, gid)
        self.tree.blockSignals(False)

    def remove_item(self, key: list[str]):
        self.tree.remove_item(key)

    def show_property(self, _: QtCore.QModelIndex = None):
        if not self.propwidget.isExpanded():
            return

        sid = []
        for i in self.tree.selectedIndexes():
            if i.column() == 0:
                sid.append(i.data(role=QtCore.Qt.ItemDataRole.UserRole))

        self.startShowProperty.emit(sid)

    @QtCore.pyqtSlot(dict)
    def set_properties(self, props: dict):
        self.proptable.populate(props)

    def change_property(self, key1, key2, value):
        if key2 == 'Value':
            try:
                value = float(value)
            except ValueError:
                QtWidgets.QMessageBox.warning(
                    self,
                    'Invalid entry',
                    f'Value {value!r} is not a valid number for `value`.')
                return

            ids = []
            for item in self.tree.selectedItems():
                ids.append(item.data(0, QtCore.Qt.ItemDataRole.UserRole))
                item.setToolTip(0, str(value))
        else:
            ids = [item.data(0, QtCore.Qt.ItemDataRole.UserRole) for item in self.tree.selectedItems()]

        self.propertyChanged.emit(ids, key1, key2, value)

    def uncheck_sets(self, sets: list[str]):
        self.tree.uncheck_sets(sets)

    def set_name(self, sid, value):
        self.tree.set_name(sid, value)

    def catchthemall(self):
        from gui_qt.lib.pokemon import QPoke

        dialog = QPoke( parent=self)
        dialog.exec()

    @property
    def management(self):
        return self.tree.management

    @management.setter
    def management(self, value):
        self.tree.management = value