Pamhyr2/src/View/Network/GraphWidget.py

# GraphWidget.py -- Pamhyr
# Copyright (C) 2023  INRAE
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.

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

import math
import logging

from tools import timer

from PyQt5.QtCore import (
    Qt, QPoint, QPointF, QSizeF, QLineF, QRectF,
    pyqtSlot, pyqtSignal, QCoreApplication,
)
from PyQt5.QtGui import (
    QPainter, QColor, QBrush, QPen, QPainterPath,
    QPolygonF,
)
from PyQt5.QtWidgets import (
    QApplication, QGraphicsScene, QGraphicsView,
    QGraphicsItem, QGraphicsTextItem, QMenu,
)

from Model.Network.Node import Node
from Model.Network.Edge import Edge
from Model.Network.Graph import Graph

from View.Network.UndoCommand import *

logger = logging.getLogger()

_translate = QCoreApplication.translate


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.setFlag(QGraphicsItem.ItemIsMovable)
        self.setFlag(QGraphicsItem.ItemSendsGeometryChanges)
        self.setCacheMode(QGraphicsItem.DeviceCoordinateCache)
        self.setZValue(1)

    def type(self):
        return NodeItem.Type

    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

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

        # Select color in function of node position in graph and
        # status
        color = Qt.darkBlue
        if self.graph.selected_new_edge_src_node() == self:
            color = Qt.darkRed
        elif self.graph.selected_item() == self:
            color = Qt.red
        elif not self.graph.graph.is_enable_node(self.node):
            color = Qt.darkGray
        elif self.graph.graph.is_upstream_node(self.node):
            color = Qt.yellow
        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.node_change_position(value, self)

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

    def mousePressEvent(self, event):
        self.update()
        if not self.graph._only_display:
            super(NodeItem, self).mousePressEvent(event)

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

    def mouseMoveEvent(self, event):
        self.update()
        if not self.graph._only_display:
            super(NodeItem, self).mouseMoveEvent(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.setAcceptedMouseButtons(Qt.NoButton)

    def type(self):
        return Edge.Type

    def boundingRect(self):
        # Rectangle of edge for display update
        pen_width = 2.0
        extra = (pen_width + 5) / 2.0

        return QRectF(
            self.src_node.pos(),
            QSizeF(
                self.dest_node.pos().x() - self.src_node.pos().x(),
                self.dest_node.pos().y() - self.src_node.pos().y()
            )
        ).normalized().adjusted(-extra, -extra, extra, extra)

    def shape(self):
        # Shape around edge for mouse selection
        vec = self.dest_node.pos() - self.src_node.pos()
        vec = vec * (5 / math.sqrt(QPointF.dotProduct(vec, vec)))
        extra = QPointF(vec.y(), -vec.x())

        result = QPainterPath(self.src_node.pos() - vec + extra)
        result.lineTo(self.src_node.pos() - vec - extra)
        result.lineTo(self.dest_node.pos() + vec - extra)
        result.lineTo(self.dest_node.pos() + vec + extra)
        result.closeSubpath()

        return result

    @timer
    def paint(self, painter, option, widget):
        # Draw shape of the edge
        # color = QColor(Qt.white)
        # color.setAlpha(128)
        # painter.setBrush(color)
        # painter.drawPath(self.shape())

        line = QLineF(self.src_node.pos(), self.dest_node.pos())
        if line.length() == 0.0:
            return

        # Select color
        color = Qt.darkBlue
        if self.graph.selected_item() == self:
            color = Qt.red
        elif self.graph.current_edge() == self:
            color = Qt.blue
        elif not self.graph.graph.is_enable_edge(self.edge):
            color = Qt.darkGray

        painter.setPen(
            QPen(
                color, 2, Qt.SolidLine, Qt.RoundCap, Qt.RoundJoin
            )
        )
        # Draw the line
        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
        brush = QBrush()
        brush.setColor(color)
        brush.setStyle(Qt.SolidPattern)

        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
        )
        poly = QPolygonF([line_center, arrow_p1, arrow_p2])
        path = QPainterPath()
        path.addPolygon(poly)

        painter.drawPolygon(poly)
        painter.fillPath(path, brush)


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.set_custom_pos(self.item.pos())
        self.setZValue(2)

    def set_custom_pos(self, pos):
        x = pos.x()
        y = pos.y() - 42  # Dont panic! The answer is 42

        self.setPos(QPointF(x, y))

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

        # Draw text background
        painter.setBrush(color)
        painter.drawRect(self.boundingRect())

        # Draw text
        super(NodeText, self).paint(painter, option, widget)

    def rename(self):
        # Update the node text
        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)

    @timer
    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,
                 min_size=(400, 400), max_size=None,
                 size=None, only_display=False, undo=None):
        super(GraphWidget, self).__init__(parent=parent)

        self.timerId = 0
        self.parent = parent
        self._state = "move"
        self._only_display = only_display
        self._undo = undo

        self.graph = graph

        self._selected_item = None
        self._selected_new_edge_src_node = None
        self._current_edge = None
        self._current_moved_node = None
        self.tmp_line = None

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

        self.m_origin_x = 0.0
        self.m_origin_y = 0.0
        self.clicked = False

        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.scale(1, 1)
        self.previousScale = 1

        if min_size:
            self.setMinimumSize(*min_size)
        if max_size:
            self.setMaximumSize(*max_size)
        if size:
            self.resize(*size)

        self.create_items()

    @timer
    def create_items(self):
        """Create all items and draw its

        Returns:
            Nothing
        """
        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)

        curr_edge = self.graph.current_reach()

        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)
            if edge == curr_edge:
                self._current_edge = iedge

            self.scene().addItem(iedge)
            self.edge_items.append(iedge)

    def state(self, status):
        """Set the current status of draw widget

        Set the current status of draw widget, like "move", "add" or
        "del"

        Args:
            status: String of current status

        Returns:
            Nothing
        """
        self._state = status

    def add_node(self, pos):
        """Create a new node in graph and draw it

        Args:
            pos: The position of new node

        Returns:
            Nothing
        """
        node = self.graph.create_node(pos.x(), pos.y())
        self._undo.push(
            AddNodeCommand(
                self.graph, node
            )
        )
        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):
        """Delete a node and update display

        Args:
            item: The node item to delete

        Returns:
            Nothing
        """
        node = item.node
        self._undo.push(
            DelNodeCommand(
                self.graph, node
            )
        )

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

    def rename_nodes(self):
        """Update all nodes item name in scene

        Returns:
            Nothing
        """
        for i in self.texts:
            if type(i) is NodeItem:
                self.texts[i].rename()

    def enable_edge(self, edge, prev):
        self._undo.push(
            EnableEdgeCommand(
                edge.edge, not prev
            )
        )
        self.changeEdge.emit(self.sender())
        self.changeNode.emit(self.sender())
        self.display_update()

    def reverse_edge(self, edge):
        self._undo.push(
            ReverseEdgeCommand(
                edge.edge
            )
        )
        self.changeEdge.emit(self.sender())
        self.display_update()

    def display_update(self):
        """Clear the scene and redraw it

        Returns:
            Nothing
        """
        self.scene().clear()
        self.create_items()

    def add_edge(self, node1, node2):
        """Create a new edge in graph and draw it

        Args:
            node1: The source node
            node2: The destination node

        Returns:
            Nothing
        """
        if node1 == node2:
            return

        edge = self.graph.create_edge(node1.node, node2.node)
        self._undo.push(
            AddEdgeCommand(
                self.graph, edge
            )
        )
        self.set_selected_item(None)
        self.set_selected_new_edge_src_node(None)
        # Reset the temporary line
        self.tmp_line = None

        # Clear all scene and redraw it
        self.scene().clear()
        self.create_items()

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

    def del_edge(self, item):
        """Delete an edge and update display

        Args:
            item: The edge item to delete

        Returns:
            Nothing
        """
        edge = item.edge
        self._undo.push(
            DelEdgeCommand(
                self.graph, edge
            )
        )

        # Clear all scene and redraw it
        self.scene().clear()
        self.create_items()

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

    def update_edges(self, node):
        """Update display of all edges linked with the node

        Args:
            node: The node item linked with edges to update

        Returns:
            Nothing
        """
        edges = list(
            filter(
                lambda ie: (ie.edge.node1 == node.node or
                            ie.edge.node2 == node.node),
                self.edge_items
            )
        )

        for edge in edges:
            edge.update()

    def node_change_position(self, pos, node):
        """Update node position and node text position

        Args:
            pos: The new position
            node: The node item

        Returns:
            Nothing
        """
        node.node.setPos(pos.x(), pos.y())
        self.texts[node].set_custom_pos(pos)
        self.texts[node].update()

    def selected_item(self):
        """Current selected item

        Returns:
            Item if item ar selected, otherelse None
        """
        return self._selected_item

    def set_selected_item(self, item):
        """Set current selected item

        Args:
            item: The new item to select

        Returns:
            Nothing
        """
        try:
            previous_item = self._selected_item
            self._selected_item = item

            if previous_item:
                previous_item.update()

                if item:
                    item.update()
        except Exception as e:
            logger.warning(str(e))

    def selected_new_edge_src_node(self):
        """The current node item selected to add new edge

        Returns:
            Item if item ar selected, otherelse None
        """
        return self._selected_new_edge_src_node

    def set_selected_new_edge_src_node(self, node):
        """Set the current node item selected to add new edge

        Args:
            node: The new node to select

        Returns:
            Nothing
        """
        previous_node = self._selected_new_edge_src_node
        self._selected_new_edge_src_node = node

        if previous_node:
            previous_node.update()

    def current_edge(self):
        """The current selected edge

        Returns:
            Item if edge are selected, otherelse None
        """
        return self._current_edge

    def set_current_edge(self, edge):
        """Set the current edge item selected

        Args:
            edge: The new edge to select

        Returns:
            Nothing
        """
        try:
            previous_edge = self._current_edge
            self._current_edge = edge
            self.graph.set_current_reach(edge.edge)

            if previous_edge:
                previous_edge.update()
        except Exception as e:
            logger.warning(str(e))

    def reset_selection(self):
        """Reset all the selected items

        Returns:
            Nothing
        """
        self.set_selected_new_edge_src_node(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()
        else:
            super(GraphWidget, self).keyPressEvent(event)

    def drawBackground(self, painter, rect):
        sceneRect = self.sceneRect()
        painter.fillRect(
            rect.intersected(sceneRect),
            QBrush(Qt.lightGray)
        )
        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

        self.scale(scaleFactor, scaleFactor)

    def mousePressEvent(self, event):
        pos = self.mapToScene(event.pos())
        self.clicked = True

        if event.buttons() & Qt.RightButton:
            self.update()
            super(GraphWidget, self).mousePressEvent(event)
            return

        # Move item and select edge item
        if self._state == "move":
            self._selected_new_edge_src_node = None

            items = self.items(event.pos())
            if items and type(items[0]) is EdgeItem:
                edge = items[0]
                if edge:
                    self.set_current_edge(edge)
            elif items and type(items[0]) is NodeItem:
                self._mouse_origin_x = pos.x()
                self._mouse_origin_y = pos.y()
                self._current_moved_node = items[0]

        # Add nodes and edges
        elif self._state == "add":
            items = self.items(event.pos())
            nodes = list(filter(lambda i: type(i) is NodeItem, items))
            if not nodes:
                self.add_node(pos)
            else:
                if self.selected_new_edge_src_node() is None:
                    self.set_selected_new_edge_src_node(nodes[0])
                else:
                    self.add_edge(self.selected_new_edge_src_node(), nodes[0])

        # Delete nodes and edges
        elif self._state == "del":
            self._selected_new_edge_src_node = None
            items = list(
                filter(
                    lambda i: type(i) is NodeItem or type(i) is EdgeItem,
                    self.items(event.pos())
                )
            )
            if len(items) > 0:
                item = items[0]
                if type(item) is NodeItem:
                    self.del_node(item)
                elif type(item) is EdgeItem:
                    self.del_edge(item)

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

    def mouseReleaseEvent(self, event):
        self.clicked = False

        if self._only_display:
            return

        if self._state == "move":
            if self._current_moved_node is not None:
                pos = self.mapToScene(event.pos())
                self._undo.push(
                    SetNodePosCommand(
                        self._current_moved_node,
                        (pos.x(), pos.y()),
                        (self._mouse_origin_x,
                         self._mouse_origin_y)
                    )
                )

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

    def mouseMoveEvent(self, event):
        pos = self.mapToScene(event.pos())

        # Selecte item on the fly
        items = self.items(event.pos())
        selectable_items = list(
            filter(
                lambda i: (type(i) is NodeItem or type(i) is EdgeItem),
                items
            )
        )

        if selectable_items:
            self.set_selected_item(selectable_items[0])
        elif not self.clicked:
            self.set_selected_item(None)

        # Update temporary line
        if self.selected_new_edge_src_node() is not None:
            if self.tmp_line is None:
                self.tmp_line = NewEdgeLine(
                    self.selected_new_edge_src_node().pos(),
                    pos
                )
                self.scene().addItem(self.tmp_line)
            else:
                self.tmp_line.dest = pos
            self.tmp_line.update()

        # If state is "move"
        if self._state == "move":
            # Move on scene
            if (not self.selected_item() and
                    event.buttons() & Qt.LeftButton):
                old_p = self.mapToScene(self.m_origin_x, self.m_origin_y)
                new_p = self.mapToScene(event.pos())
                translation = new_p - old_p

                self.translate(translation.x(), translation.y())

                self.m_origin_x = event.x()
                self.m_origin_y = event.y()

            # Propagate event
            self.update()
            super(GraphWidget, self).mouseMoveEvent(event)

    # Contextual menu

    def _menu_default(self, event, pos, items, menu):
        add_node = menu.addAction(_translate("Network", "Add node"))

        action = menu.exec_(self.mapToGlobal(event.pos()))

        if action == add_node:
            self.add_node(pos)

    def _menu_node(self, event, pos, items, menu):
        delete = menu.addAction(_translate("Network", "Delete the node"))
        disable = menu.addAction(_translate("Network", "Disable the node"))

        action = menu.exec_(self.mapToGlobal(event.pos()))

        if action == delete:
            self.del_node(items[0])

    def _menu_edge(self, event, pos, items, menu):
        delete = menu.addAction(_translate("Network", "Delete the reach"))

        if self.graph.is_enable_edge(items[0].edge):
            enable = menu.addAction(_translate("Network", "Disable the reach"))
            is_enable = True
        else:
            enable = menu.addAction(_translate("Network", "Enable the reach"))
            is_enable = False

        reverse = menu.addAction(_translate(
            "Network", "Reverse the reach orientation"))

        action = menu.exec_(self.mapToGlobal(event.pos()))

        if action == delete:
            self.del_edge(items[0])
        elif action == enable:
            self.enable_edge(items[0], is_enable)
        elif action == reverse:
            self.reverse_edge(items[0])

    def contextMenuEvent(self, event):
        if self._only_display:
            return

        pos = self.mapToScene(event.pos())
        items = self.items(event.pos())

        menu = QMenu(self)
        if len(items) == 0:
            self._menu_default(event, pos, items, menu)
        elif type(items[0]) is NodeItem:
            self._menu_node(event, pos, items, menu)
        elif type(items[0]) is EdgeItem:
            self._menu_edge(event, pos, items, menu)