GraphWidget.py

# -*- coding: utf-8 -*-

import math

from model.network.Node import Node
from model.network.Edge import Edge
from model.network.Graph import Graph

from PyQt5.QtCore import (
    qAbs, QLineF, QPointF, QPoint, qrand, QRectF, QSizeF, qsrand,
    Qt, QTime, pyqtSlot, pyqtSignal,
)
from PyQt5.QtGui import (
    QBrush, QColor, QLinearGradient, QPainter,
    QPainterPath, QPen, QPolygonF, QRadialGradient,
    QFont,
)
from PyQt5.QtWidgets import (
    QApplication, QGraphicsItem, QGraphicsScene,
    QGraphicsView, QStyle, QGraphicsTextItem,
)

class LandMark(QGraphicsItem):
    def __init__(self, width, height, parent=None):
        super(LandMark, self).__init__()

        self.width = width
        self.height = height

        self.center = QPoint(width / 2, height / 2)
        self.parent = parent

    def paint(self, painter, option, widget):
        painter.setPen(Qt.NoPen)
        painter.setPen(QPen(Qt.black, 0))

        painter.drawLine(
            -10 + self.center.x(), 0 + self.center.y(),
            10 + self.center.x(), 0 + self.center.y()
        )
        painter.drawLine(
            0 + self.center.x(), -10 + self.center.y(),
            0 + self.center.x(), 10 + self.center.y()
        )

    def boundingRect(self):
        return QRectF(
            -10 + self.center.x(), -10 + self.center.y(),
            10 + self.center.x(), 10 + self.center.y()
        )


class NodeItem(QGraphicsItem):
    Type = QGraphicsItem.UserType + 1

    def __init__(self, node, graph_widget):
        super(NodeItem, self).__init__()

        self.node = node
        self.setPos(QPointF(self.node.pos.x, self.node.pos.y))

        self.graph = graph_widget
        self.newPos = QPointF()

        self.setFlag(QGraphicsItem.ItemIsMovable)
        self.setFlag(QGraphicsItem.ItemSendsGeometryChanges)
        self.setCacheMode(QGraphicsItem.DeviceCoordinateCache)
        self.setZValue(1)

    def type(self):
        return NodeItem.Type

    def advance(self):
        if self.newPos == self.pos():
            return False

        self.setPos(self.newPos)
        return True

    def boundingRect(self):
        adjust = 2.0
        return QRectF(-10 - adjust, -10 - adjust, 20 + adjust, 20 + adjust)

    def shape(self):
        path = QPainterPath()
        path.addEllipse(-10, -10, 20, 20)
        return path

    def paint(self, painter, option, widget):
        painter.setPen(Qt.NoPen)

        color = Qt.yellow
        if self.graph.selectedNewEdgeSrcNode() == self:
            color = Qt.darkRed
        elif self.graph.selectedItem() == self:
            color = Qt.red
        elif self.graph.graph.is_upstream_node(self.node):
            color = Qt.blue
        elif self.graph.graph.is_downstream_node(self.node):
            color = Qt.green

        painter.setBrush(QBrush(color))
        painter.drawEllipse(-10, -10, 20, 20)

    def itemChange(self, change, value):
        if change == QGraphicsItem.ItemPositionHasChanged:
            self.graph.itemMoved()
            self.graph.nodeChangePosition(value, self)

        self.graph.update_edges(self)
        return super(NodeItem, self).itemChange(change, value)

    def mousePressEvent(self, event):
        self.update()
        super(NodeItem, self).mousePressEvent(event)

    def mouseReleaseEvent(self, event):
        self.update()
        super(NodeItem, self).mouseReleaseEvent(event)

class EdgeItem(QGraphicsItem):
    Type = QGraphicsItem.UserType + 2

    def __init__(self, src_node_item, dest_node_item, edge, graph_widget):
        super(EdgeItem, self).__init__()

        self.src_node = src_node_item
        self.dest_node = dest_node_item
        self.edge = edge

        self.graph = graph_widget

        self.src_pos = src_node_item.pos()
        self.dest_pos = dest_node_item.pos()

        self.setAcceptedMouseButtons(Qt.NoButton)
        self.adjust()

    def type(self):
        return Edge.Type

    def sourceNode(self):
        return self.src_node.node

    def setSourceNode(self, node_item):
        self.src_node = node_item
        self.adjust()

    def destNode(self):
        return self.dest

    def setDestNode(self, node_item):
        self.dest_node = node_item
        self.adjust()

    def adjust(self):
        line = QLineF(
            self.mapFromItem(self.src_node, 0, 0),
            self.mapFromItem(self.dest_node, 0, 0)
        )
        length = line.length()

        self.prepareGeometryChange()

    def boundingRect(self):
        pen_width = 3.0
        extra = (pen_width + 5) / 2.0

        return QRectF(
            self.src_pos,
            QSizeF(
                self.dest_pos.x() - self.src_pos.x(),
                self.dest_pos.y() - self.src_pos.y()
            )
        ).normalized().adjusted(-extra, -extra, extra, extra)

    def paint(self, painter, option, widget):
        # Draw the line itself.
        line = QLineF(self.src_node.pos(), self.dest_node.pos())

        if line.length() == 0.0:
            return

        color = Qt.blue
        if self.graph.selectedItem() == self:
            color = Qt.red

        painter.setPen(QPen(color, 3, Qt.SolidLine, Qt.RoundCap,
                            Qt.RoundJoin))
        painter.drawLine(line)

        line_center = QPointF(
            (self.src_node.pos().x() + self.dest_node.pos().x()) / 2,
            (self.src_node.pos().y() + self.dest_node.pos().y()) / 2,
        )

        # Draw the arrows if there's enough room.
        angle = math.acos(line.dx() / line.length())
        if line.dy() >= 0:
            angle = (math.pi * 2.0) - angle

        size = 10.0
        arrow_p1 = line_center + QPointF(
            math.sin(angle - math.pi / 3) * size,
            math.cos(angle - math.pi / 3) * size
        )
        arrow_p2 = line_center + QPointF(
            math.sin(angle - math.pi + math.pi / 3) * size,
            math.cos(angle - math.pi + math.pi / 3) * size
        )

        painter.drawPolygon(QPolygonF([line_center, arrow_p1, arrow_p2]))

class NodeText(QGraphicsTextItem):
    def __init__(self, node_item):
        super(NodeText, self).__init__()

        self.item = node_item
        self.setPlainText(self.item.node.name)
        self.setDefaultTextColor(Qt.black)
        #self.setPos(self.item.pos())
        self.set_custom_pos(self.item.pos())
        self.setZValue(2)

    def set_custom_pos(self, pos):
        x = pos.x()
        y = pos.y() - 42

        self.setPos(QPointF(x,y))


    def paint(self, painter, option, widget):
        color = QColor(Qt.white)
        color.setAlpha(128)

        painter.setBrush(color)
        painter.drawRect(self.boundingRect())

        super(NodeText, self).paint(painter, option, widget)

    def rename(self):
        self.setPlainText(self.item.node.name)

class NewEdgeLine(QGraphicsItem):
    def __init__(self, src, dest):
        super(NewEdgeLine, self).__init__()

        self.src = src
        self.dest = dest

    def boundingRect(self):
        pen_width = 2.0
        extra = (pen_width + 5) / 2.0

        return QRectF(
            self.src,
            QSizeF(
                self.dest.x() - self.src.x(),
                self.dest.y() - self.src.y()
            )
        ).normalized().adjusted(-extra, -extra, extra, extra)

    def paint(self, painter, option, widget):
        line = QLineF(self.src, self.dest)
        if line.length() == 0.0:
            return

        color = Qt.darkGray
        painter.setPen(QPen(color, 2, Qt.SolidLine, Qt.RoundCap,
                            Qt.RoundJoin))
        painter.drawLine(line)

class GraphWidget(QGraphicsView):
    changeEdge = pyqtSignal(object)
    changeNode = pyqtSignal(object)

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

        self.timerId = 0
        self.parent = parent
        self._state = "move"

        self._selectedItem = None
        self._selectedNewEdgeSrcNode = None
        self.tmp_line = None

        self.node_items = []
        self.edge_items = []
        self.texts = {}

        scene = QGraphicsScene(self)
        scene.setItemIndexMethod(QGraphicsScene.NoIndex)
        scene.setSceneRect(0, 0, 2000, 2000)
        self.setScene(scene)
        self.setCacheMode(QGraphicsView.CacheBackground)
        self.setViewportUpdateMode(QGraphicsView.BoundingRectViewportUpdate)
        self.setRenderHint(QPainter.Antialiasing)
        self.setTransformationAnchor(QGraphicsView.AnchorUnderMouse)
        self.setResizeAnchor(QGraphicsView.AnchorViewCenter)

        self.graph = graph

        self.scale(1, 1)
        self.previousScale = 1
        self.setMinimumSize(400, 400)

        self.create_items()

    def create_items(self):
        self.node_items = []
        self.edge_items = []
        self.texts = {}

        for node in self.graph.nodes():
            inode = NodeItem(node, self)
            self.texts[inode] = NodeText(inode)
            self.scene().addItem(self.texts[inode])
            self.scene().addItem(inode)
            self.node_items.append(inode)

        for edge in self.graph.edges():
            n1 = list(
                filter(
                    lambda n: n.node.name == edge.node1.name,
                    self.node_items
                )
            )

            n2 = list(
                filter(
                    lambda n: n.node.name == edge.node2.name,
                    self.node_items
                )
            )

            iedge = EdgeItem(n1[0], n2[0], edge, self)
            self.scene().addItem(iedge)
            self.edge_items.append(iedge)

        # self.mark = LandMark(1000, 1000, self)
        # self.scene().addItem(self.mark)

    def state(self, status):
        self._state = status

    def add_node(self, pos):
        node = self.graph.add_node(pos.x(), pos.y())
        inode = NodeItem(node, self)
        self.scene().addItem(inode)
        self.node_items.append(inode)
        self.texts[inode] = NodeText(inode)
        self.scene().addItem(self.texts[inode])

        self.changeNode.emit(self.sender())

    def del_node(self, item):
        node = item.node
        edges = list(
            filter(
                lambda ie: ie.edge.node1 == node or ie.edge.node2 == node,
                self.edge_items
            )
        )

        self.setSelectedItem(None)

        for edge in edges:
            self.graph.remove_edge(edge.edge.name)

        self.graph.remove_node(node.name)

        self.scene().clear()
        self.create_items()
        self.changeNode.emit(self.sender())
        self.changeEdge.emit(self.sender())

    def rename_nodes(self):
        for i in self.texts:
            if type(i) == NodeItem:
                self.texts[i].rename()

    def add_edge(self, node1, node2):
        if node1 == node2:
            return

        edge = self.graph.add_edge(node1.node, node2.node)
        # iedge = EdgeItem(node1, node2, edge, self)
        # self.scene().addItem(iedge)
        # self.edge_items.append(iedge)
        self.setSelectedItem(None)
        self.setSelectedNewEdgeSrcNode(None)

        self.tmp_line = None
        self.scene().clear()
        self.create_items()

        self.changeEdge.emit(self.sender())

    def del_edge(self, item):
        edge = item.edge
        self.graph.remove_edge(edge.name)

        self.scene().clear()
        self.create_items()
        self.changeNode.emit(self.sender())
        self.changeEdge.emit(self.sender())


    def update_edges(self, node):
        edges = list(
            filter(
                lambda ie: ie.edge.node1 == node.node or ie.edge.node2 == node.node,
                self.edge_items
            )
        )

        for edge in edges:
            # edge.node1.update()
            # edge.node2.update()
            edge.update()

    def itemMoved(self):
        if not self.timerId:
            self.timerId = self.startTimer(1000 / 25)

    def nodeChangePosition(self, pos, node):
        node.node.setPos(pos.x(), pos.y())
        self.texts[node].set_custom_pos(pos)
        self.texts[node].update()

    def selectedItem(self):
        return self._selectedItem

    def setSelectedItem(self, item):
        previous_item = self._selectedItem
        self._selectedItem = item

        if previous_item:
            previous_item.update()

        if item:
            item.update()

    def selectedNewEdgeSrcNode(self):
        return self._selectedNewEdgeSrcNode

    def setSelectedNewEdgeSrcNode(self, node):
        previous_node = self._selectedNewEdgeSrcNode
        self._selectedNewEdgeSrcNode = node

        if previous_node:
            previous_node.update()

    def reset_selection(self):
        self.setSelectedNewEdgeSrcNode(None)
        if self.tmp_line is not None:
            self.tmp_line = None
            self.scene().clear()
            self.create_items()

    def keyPressEvent(self, event):
        key = event.key()

        if key == Qt.Key_Plus:
            self.scaleView(1.2)
        elif key == Qt.Key_Minus:
            self.scaleView(1 / 1.2)
        elif key == Qt.Key_Escape:
            self.reset_selection()
        elif key == Qt.Key_Space or key == Qt.Key_Enter:
            for item in self.scene().items():
                if isinstance(item, Node):
                    item.setPos(-150 + qrand() % 300, -150 + qrand() % 300)
        else:
            super(GraphWidget, self).keyPressEvent(event)

    def drawBackground(self, painter, rect):
        sceneRect = self.sceneRect()

        gradient = QLinearGradient(sceneRect.topLeft(), sceneRect.bottomRight())
        gradient.setColorAt(0, Qt.lightGray)
        gradient.setColorAt(1, Qt.lightGray)
        painter.fillRect(rect.intersected(sceneRect), QBrush(gradient))
        painter.setBrush(Qt.NoBrush)
        painter.drawRect(sceneRect)

    def wheelEvent(self, event):
        self.scaleView(math.pow(2.0, -event.angleDelta().y() / 240.0))

    def scaleView(self, scaleFactor):
        factor = self.transform().scale(scaleFactor, scaleFactor).mapRect(QRectF(0, 0, 1, 1)).width()

        if factor < 0.05 or factor > 1.5:
            return

        # print(f"{scaleFactor} : {factor}")

        self.scale(scaleFactor, scaleFactor)

    def mousePressEvent(self, event):
        pos = self.mapToScene(event.pos())
        if self._state == "move":
            self.mouse_origin_x = pos.x()
            self.mouse_origin_y = pos.y()

        elif self._state == "add":
            items = self.items(event.pos())
            nodes = list(filter(lambda i: type(i) == NodeItem, items))
            if not nodes:
                self.add_node(pos)
            else:
                if self.selectedNewEdgeSrcNode() is None:
                    self.setSelectedNewEdgeSrcNode(nodes[0])
                else:
                    self.add_edge(self.selectedNewEdgeSrcNode(), nodes[0])

        elif self._state == "del":
            items = list(
                filter(
                    lambda i: type(i) == NodeItem or type(i) == EdgeItem,
                    self.items(event.pos())
                )
            )
            if len(items) > 0:
                item = items[0]
                if type(item) == NodeItem:
                    self.del_node(item)
                elif type(item) == EdgeItem:
                    self.del_edge(item)

        self.update()
        super(GraphWidget, self).mousePressEvent(event)

    def mouseReleaseEvent(self, event):
        self.update()
        super(GraphWidget, self).mouseReleaseEvent(event)

    def mouseMoveEvent(self, event):
        pos = self.mapToScene(event.pos())
        items = self.items(event.pos())
        selectable_items = list(
            filter(
                lambda i: type(i) == NodeItem or type(i) == EdgeItem,
                items
            )
        )

        if selectable_items:
            self.setSelectedItem(selectable_items[0])
        else:
            self.setSelectedItem(None)

        if self.selectedNewEdgeSrcNode() is not None:
            if self.tmp_line is None:
                self.tmp_line = NewEdgeLine(
                    self.selectedNewEdgeSrcNode().pos(),
                    pos
                )
                self.scene().addItem(self.tmp_line)
            else:
                self.tmp_line.dest = pos
            self.tmp_line.update()

        if self._state == "move":
            self.update()
            super(GraphWidget, self).mouseMoveEvent(event)