modelverse/sketchUI/exm_mainwindow.py

from PyQt5.QtWidgets import QMainWindow, QAction, QActionGroup, QGraphicsItem, QGraphicsView,\
    QTableWidgetItem, QInputDialog
from PyQt5.QtGui import QIcon
from PyQt5.QtCore import QStateMachine, QState, Qt
from PyQt5.Qt import QApplication
from sketchUI.ui import Ui_MainWindow
from sketchUI.exm_scene import SketchScene, Mode
from sketchUI import mvops
from sketchUI.graphics_node_item import GraphicsNodeItem
from sketchUI.graphics_edge_item import GraphicsEdgeItem
from wrappers.modelverse import element_list_nice
import commons
from evolution.node_ops import NodeDelete, NodeAdd
from evolution.attribute_ops import AttributeDelete


class EXMMainWindow(QMainWindow, Ui_MainWindow):
    def __init__(self, model):
        QMainWindow.__init__(self)
        self.setupUi(self)
        self.setWindowTitle(model)
        self.menuActionVerify.setDisabled(True)
        self._cur_model = model

        self._scene = SketchScene(model, self)
        self._scene.setSceneRect(0, 0, 200, 200)
        self.graphicsView.setScene(self._scene)
        self._scene.set_mode(Mode.SELECT)

        self.populate_types()
        self.listWidget.itemDoubleClicked.connect(self._on_list_item_clicked)

        self._setup_toolbar()
        self._setup_statemachine()

        self.menuActionDeleteModel.triggered.connect(self._on_delete_model_clicked)

        # load model
        self._load_model()

        # setup log viewer
        self.plainTextEdit.setReadOnly(True)

        # setup table view for attributes
        self.tableWidget.setColumnCount(2)
        self.tableWidget.setHorizontalHeaderLabels(["Key", "Value"])
        self.tableWidget.horizontalHeader().setStretchLastSection(True)
        self.tableWidget.itemChanged.connect(self._on_attribute_edited)

        # lastly, run state machine
        self._statemachine.start()

    def populate_types(self):
        # sync example model types to list widget
        self.listWidget.clear()
        self.listWidget.addItems(commons.get_available_types())

    def _setup_toolbar(self):
        # create action buttons on toolbar
        self.select_action = QAction("Select", self)
        self.select_action.setIcon(QIcon("sketchUI/icons/select.png"))
        self.select_action.setCheckable(True)
        self.select_action.setChecked(True)

        self.connect_action = QAction("Connect", self)
        self.connect_action.setIcon(QIcon("sketchUI/icons/connect.png"))
        self.connect_action.setCheckable(True)

        self.line_action = QAction("Draw line", self)
        self.line_action.setIcon(QIcon("sketchUI/icons/line.png"))
        self.line_action.setCheckable(True)

        self.rect_action = QAction("Draw Rectangle", self)
        self.rect_action.setIcon(QIcon("sketchUI/icons/rect.png"))
        self.rect_action.setCheckable(True)

        self.circle_action = QAction("Draw Circle", self)
        self.circle_action.setIcon(QIcon("sketchUI/icons/circle.png"))
        self.circle_action.setCheckable(True)

        self.pen_action = QAction("Pen", self)
        self.pen_action.setIcon(QIcon("sketchUI/icons/pen.svg"))
        self.pen_action.setCheckable(True)
        self.pen_action.setEnabled(False)

        action_group = QActionGroup(self)
        action_group.setExclusive(True)
        action_group.addAction(self.select_action)
        action_group.addAction(self.connect_action)
        action_group.addAction(self.line_action)
        action_group.addAction(self.rect_action)
        action_group.addAction(self.circle_action)
        action_group.addAction(self.pen_action)

        for item in action_group.actions():
            self.toolBar.addAction(item)

    def _setup_statemachine(self):
        self._statemachine = QStateMachine()
        state_select = QState()
        state_connect = QState()
        state_draw_line = QState()
        state_draw_rect = QState()
        state_draw_circle = QState()
        state_draw_free = QState()

        # from select
        state_select.addTransition(self.select_action.triggered, state_select)
        state_select.addTransition(self.connect_action.triggered, state_connect)
        state_select.addTransition(self.line_action.triggered, state_draw_line)
        state_select.addTransition(self.rect_action.triggered, state_draw_rect)
        state_select.addTransition(self.circle_action.triggered, state_draw_circle)
        state_select.addTransition(self.pen_action.triggered, state_draw_free)
        state_select.entered.connect(self._state_select_entered)
        # from connect
        state_connect.addTransition(self.connect_action.triggered, state_connect)
        state_connect.addTransition(self.select_action.triggered, state_select)
        state_connect.addTransition(self.line_action.triggered, state_draw_line)
        state_connect.addTransition(self.rect_action.triggered, state_draw_rect)
        state_connect.addTransition(self.circle_action.triggered, state_draw_circle)
        state_connect.addTransition(self.pen_action.triggered, state_draw_free)
        state_connect.entered.connect(self._state_connect_entered)
        # from rectangle
        state_draw_rect.addTransition(self.rect_action.triggered, state_draw_rect)
        state_draw_rect.addTransition(self.connect_action.triggered, state_connect)
        state_draw_rect.addTransition(self.select_action.triggered, state_select)
        state_draw_rect.addTransition(self.line_action.triggered, state_draw_line)
        state_draw_rect.addTransition(self.circle_action.triggered, state_draw_circle)
        state_draw_rect.addTransition(self.pen_action.triggered, state_draw_free)
        state_draw_rect.entered.connect(self._state_draw_rect_entered)
        # from line
        state_draw_line.addTransition(self.line_action.triggered, state_draw_line)
        state_draw_line.addTransition(self.rect_action.triggered, state_draw_rect)
        state_draw_line.addTransition(self.connect_action.triggered, state_connect)
        state_draw_line.addTransition(self.select_action.triggered, state_select)
        state_draw_line.addTransition(self.circle_action.triggered, state_draw_circle)
        state_draw_line.addTransition(self.pen_action.triggered, state_draw_free)
        state_draw_line.entered.connect(self._state_draw_line_entered)
        # from circle
        state_draw_circle.addTransition(self.circle_action.triggered, state_draw_circle)
        state_draw_circle.addTransition(self.line_action.triggered, state_draw_line)
        state_draw_circle.addTransition(self.rect_action.triggered, state_draw_rect)
        state_draw_circle.addTransition(self.connect_action.triggered, state_connect)
        state_draw_circle.addTransition(self.select_action.triggered, state_select)
        state_draw_circle.addTransition(self.pen_action.triggered, state_draw_free)
        state_draw_circle.entered.connect(self._state_draw_circle_entered)
        # from freehand
        state_draw_free.addTransition(self.pen_action.triggered, state_draw_free)
        state_draw_free.addTransition(self.circle_action.triggered, state_draw_circle)
        state_draw_free.addTransition(self.line_action.triggered, state_draw_line)
        state_draw_free.addTransition(self.rect_action.triggered, state_draw_rect)
        state_draw_free.addTransition(self.connect_action.triggered, state_connect)
        state_draw_free.addTransition(self.select_action.triggered, state_select)
        state_draw_free.entered.connect(self._state_draw_free_entered)

        self._statemachine.addState(state_select)
        self._statemachine.addState(state_connect)
        self._statemachine.addState(state_draw_rect)
        self._statemachine.addState(state_draw_line)
        self._statemachine.addState(state_draw_circle)
        self._statemachine.addState(state_draw_free)
        self._statemachine.setInitialState(state_select)

    def _state_select_entered(self):
        self._scene.set_mode(Mode.SELECT)
        self._make_items_movable(True)
        self._enable_box_select(True)

    def _state_connect_entered(self):
        self._scene.set_mode(Mode.CONNECT)
        self._make_items_movable(False)
        self._enable_box_select(False)

    def _state_draw_rect_entered(self):
        self._scene.set_mode(Mode.RECT)
        self._make_items_movable(False)
        self._enable_box_select(False)

    def _state_draw_line_entered(self):
        self._scene.set_mode(Mode.LINE)
        self._make_items_movable(False)
        self._enable_box_select(False)

    def _state_draw_circle_entered(self):
        self._scene.set_mode(Mode.CIRCLE)
        self._make_items_movable(False)
        self._enable_box_select(False)

    def _state_draw_free_entered(self):
        self._scene.set_mode(Mode.FREE)
        self._make_items_movable(False)
        self._enable_box_select(False)

    def _make_items_movable(self, movable):
        for item in self._scene.items():
            if isinstance(item, GraphicsEdgeItem):
                # edges are never movable, but selectable
                item.setFlag(QGraphicsItem.ItemIsMovable, False)
                item.setFlag(QGraphicsItem.ItemIsSelectable, True)
            else:
                item.setFlag(QGraphicsItem.ItemIsMovable, movable)
                item.setFlag(QGraphicsItem.ItemIsSelectable, movable)

    def _enable_box_select(self, enable):
        if enable:
            self.graphicsView.setDragMode(QGraphicsView.RubberBandDrag)
        else:
            self.graphicsView.setDragMode(QGraphicsView.NoDrag)

    def _load_model(self):
        model = element_list_nice(self._cur_model)
        if not model:
            # empty
            return

        self._scene.clear()

        x_pos = -150
        y_pos = 70
        for item in model:
            typ = item["type"]
            if typ == "Node":
                # first, draw all nodes
                node_type = item["typeID"]
                self._add_node_to_scene(item["id"], node_type, x=x_pos, y=y_pos)
                x_pos += 80
                y_pos *= -1

        for item in model:
            typ = item["type"]
            if typ == "Edge":
                target = item["__target"]
                src = item["__source"]
                self._add_edge_to_scene(src, target, item["id"])

    def _add_node_to_scene(self, node_id, node_type, x=0, y=0):
        consyn = mvops.get_consyn_of(node_type)
        item = GraphicsNodeItem(node_id, node_type, consyn)
        item.setPos(x, y)
        item.setFlag(QGraphicsItem.ItemIsMovable, True)
        item.setFlag(QGraphicsItem.ItemIsSelectable, True)
        self._scene.addItem(item)

    def _add_edge_to_scene(self, from_id, to_id, edge_id):
        from_item = None
        to_item = None
        for item in self._scene.items():
            try:
                node_id = item.node_id
            except AttributeError:
                # no node item, continue
                continue
            if node_id == from_id:
                from_item = item
                continue
            if node_id == to_id:
                to_item = item

        self._scene.draw_edge(from_item, to_item, edge_id=edge_id)

    def _on_list_item_clicked(self, event):
        # add node to model
        self.plainTextEdit.appendPlainText("Adding node of type {} to model".format(event.text()))
        add_handler = NodeAdd()

        # ask if global or local add node
        scope, ok = QInputDialog.getItem(self.parent(), "Select scope", "Scope", ["Local", "Global"])
        if not ok:
            return

        if scope == "Global":
            add_handler.execute(self._cur_model, event.text(), local=False)
        else:
            add_handler.execute(self._cur_model, event.text(), local=True)
        # repair any possibly invalidated instance models
        add_handler.repair()
        # render node by reloading the whole model (since we can not get the node id of the newly added model)
        self._load_model()

    def _on_attribute_edited(self, item):
        # type: (QTableWidgetItem) -> None
        """
        An attribute value was edited.
        If the new entered value is empty, delete the attribute.
        """
        row = self.tableWidget.row(item)
        attr_key = self.tableWidget.item(row, 0).text()
        attr_val = self.tableWidget.item(row, 1).text()
        node = self._scene.selectedItems()[0]

        if not attr_val:
            # ask if global or local delete
            scope, ok = QInputDialog.getItem(self.parent(), "Select scope", "Scope", ["Local", "Global"])
            if not ok:
                return
            del_handler = AttributeDelete()
            self.plainTextEdit.appendPlainText("Deleting attribute {}".format(attr_key))

            if scope == "Global":
                del_handler.execute(self._cur_model, node.node_id, attr_key, local=False)
            else:
                del_handler.execute(self._cur_model, node.node_id, attr_key, local=True, check_if_last=True)

            # update view
            self.tableWidget.removeRow(row)
        else:
            self.plainTextEdit.appendPlainText("Updating value of attribute {} to {}".format(attr_key, attr_val))
            mvops.update_attribute_val(self._cur_model, node.node_id, attr_key, attr_val)

    def add_new_attribute(self, key, val="unknown"):
        """
        Adds a new attribute to the view with key "key" and optional val.
        """
        selected_node = self._scene.selectedItems()[0]
        self.plainTextEdit.appendPlainText("Adding new attribute with key {} to node {}".format(key, selected_node.get_type()))
        self.tableWidget.blockSignals(True)
        table_item_key = QTableWidgetItem(key)
        table_item_key.setFlags(Qt.ItemIsSelectable | Qt.ItemIsEnabled)
        table_item_val = QTableWidgetItem(val)
        cur_row_cnt = self.tableWidget.rowCount()
        self.tableWidget.insertRow(cur_row_cnt)
        self.tableWidget.setItem(cur_row_cnt, 0, table_item_key)
        self.tableWidget.setItem(cur_row_cnt, 1, table_item_val)
        self.tableWidget.blockSignals(False)

    def _on_delete_model_clicked(self, event):
        # delete the open example model by iteratively deleting all nodes
        self.plainTextEdit.appendPlainText("Deleting model ...")
        self.plainTextEdit.repaint()

        QApplication.setOverrideCursor(Qt.WaitCursor)
        delhander = NodeDelete()
        for item in self._scene.items():
            if not isinstance(item, GraphicsNodeItem):
                continue
            self.plainTextEdit.appendPlainText("Deleting node {}".format(item.node_id))
            self.plainTextEdit.repaint()
            delhander.execute(self._cur_model, item.node_id, local=True, check_if_last=True)

        # delete the empty model from the Modelverse
        if mvops.delete_example_model(self._cur_model):
            self.close()
        else:
            self.plainTextEdit.appendPlainText("Error: Delete failed")
            QApplication.restoreOverrideCursor()