StoryCheckTools/graph/directed_acyclic_graph/DAGPresent.py

from graph.directed_acyclic_graph.DAGLayout import DAGGraph, DAGOrderHelper
from graph.DataType import Arrow, Point
from typing import List, Dict
from PyQt5.QtWidgets import QWidget, QApplication, QGraphicsView, QGraphicsScene, QGraphicsItem
from PyQt5.QtCore import QRectF, Qt, QPointF
from PyQt5.QtGui import QFont, QBrush, QPen, QPainter, QPainterPath
import sys
from enum import Enum
from abc import abstractmethod


class GraphNode:
    @abstractmethod
    def boundingRect(self) -> QRectF:
        raise NotImplementedError("boundingRect")

    @abstractmethod
    def pos(self) -> QPointF:
        raise NotImplementedError("pos")


class StoryNodeType(Enum):
    """
    故事节点类型
    """
    StoryStartNode = 0,
    FragmentNode = 1,


class ActivePresentNode(QGraphicsItem, GraphNode):
    """
    故事情节名称节点
    """

    def __init__(self, name: str, type: StoryNodeType, font: QFont, parent):
        QGraphicsItem.__init__(self, parent)

        self.node_type_within_story = type
        self.fragment_name = name
        self.font_bind = font
        self.data_bind: object = None
        pass

    def boundingRect(self) -> QRectF:
        size = self.font_bind.pixelSize()
        if self.node_type_within_story == StoryNodeType.FragmentNode:
            return QRectF(0, 0, size * len(self.fragment_name) + 10, size + 10)
        else:
            return QRectF(0, 0, size * len(self.fragment_name) + 10, size + 10)

    def paint(self, painter, option, widget=...):
        outline = self.boundingRect()
        text_rect = QRectF(outline.x() + 5, outline.y() + 5, outline.width() - 10, outline.height() - 10)

        painter.save()
        if self.node_type_within_story == StoryNodeType.FragmentNode:
            painter.fillRect(outline, QBrush(Qt.gray))
        else:
            painter.fillRect(outline, Qt.green)
        painter.drawRect(outline)
        painter.setFont(self.font_bind)
        painter.drawText(text_rect, self.fragment_name)
        painter.restore()

        pass

    # 添加自定义功能函数
    def attach_data(self, data_inst: object):
        self.data_bind = data_inst
        pass

    def get_data(self) -> object:
        return self.data_bind

    pass


class PenetrateNode(QGraphicsItem, GraphNode):
    """
    贯穿连线节点
    """

    def __init__(self, width: float, parent):
        QGraphicsItem.__init__(self, parent)
        self.width_store = width
        self.data_bind: object = None
        pass

    def resize_width(self, width: float):
        self.width_store = width
        self.update()
        pass

    def boundingRect(self):
        return QRectF(0, 0, self.width_store, 8)

    def paint(self, painter, option, widget=...):
        outline = self.boundingRect()

        painter.save()
        painter.fillRect(QRectF(0, 2, outline.width(), 4), Qt.black)
        painter.restore()
        pass

    # 添加自定义功能函数
    def attach_data(self, data_inst: object):
        self.data_bind = data_inst
        pass

    def get_data(self) -> object:
        return self.data_bind


class TransitionCurve(QGraphicsItem):
    def __init__(self, start: GraphNode, end: GraphNode, prev_layrer_width:float, parent):
        QGraphicsItem.__init__(self, parent)
        self.start_node = start
        self.end_node = end
        self.prev_layer_w = prev_layrer_width
        self.outline = QRectF()
        self.data_bind: object = None
        pass

    def layout_refresh(self):
        orect = self.start_node.boundingRect()
        erect = self.end_node.boundingRect()

        xpos = self.start_node.pos().x() + self.start_node.boundingRect().width()
        width_value = self.end_node.pos().x() - self.start_node.pos().x() - orect.width()

        ypos = min(self.start_node.pos().y(), self.end_node.pos().y())
        bottom_y = max(self.start_node.pos().y() + orect.height(), self.end_node.pos().y() + erect.height())

        self.setPos(xpos, ypos)
        self.outline = QRectF(0, 0, width_value, bottom_y - ypos)
        self.update()

    def boundingRect(self):
        return self.outline
        pass

    def paint(self, painter, option, widget=...):
        outline = self.outline

        start_rect = self.start_node.boundingRect()
        end_rect = self.end_node.boundingRect()
        aj_start_pos = self.start_node.pos() + QPointF(start_rect.width(), start_rect.height()/2)
        aj_end_pos = self.end_node.pos() + QPointF(0, end_rect.height()/2)
        line_span = self.prev_layer_w - start_rect.width()

        painter.save()
        painter.setRenderHint(QPainter.RenderHint.Antialiasing)

        #if aj_start_pos.y() < aj_end_pos.y():
        start_pos = aj_start_pos - self.pos()
        end_pos = aj_end_pos - self.pos()

        line_epos = start_pos + QPointF(line_span, 0)
        control_pos0 = line_epos + QPointF((outline.width() - line_span)/3, 0)
        control_pos1 = end_pos - QPointF((outline.width() - line_span)/3, 0)

        npen = QPen(Qt.black)
        npen.setWidthF(4)
        painter.setPen(npen)

        painter.drawLine(start_pos, line_epos)

        path0 = QPainterPath()
        path0.moveTo(line_epos)
        path0.cubicTo(control_pos0, control_pos1, end_pos)
        painter.drawPath(path0)

        painter.restore()
        pass

    # 添加自定义功能函数
    def attach_data(self, data_inst: object):
        self.data_bind = data_inst
        pass

    def get_data(self) -> object:
        return self.data_bind


class Direction(Enum):
    RankLR = 0,
    RankTB = 1,


class DAGActiveView(QGraphicsView):
    def __init__(self, orie: Direction, parent):
        QGraphicsView.__init__(self, parent)
        self.setViewportUpdateMode(QGraphicsView.ViewportUpdateMode.FullViewportUpdate)

        self.layer_span = 200
        self.node_span = 20

        self.orientation = orie
        self.scene_bind = QGraphicsScene(self)
        self.setScene(self.scene_bind)

        font = QFont()
        font.setPixelSize(20)
        self.setFont(font)

        pass

    def update_with_edges(self, arrows: List[Arrow]) -> None:
        tools = DAGGraph()
        tools.rebuild_from_edges(arrows)
        tools.graph_layout()

        total_nodes = tools.nodes_with_layout
        previous_node_end = 0
        previous_graphics_nodes = []
        # 迭代呈现层
        for layer_idx in range(0, tools.max_layer_count):
            current_level_nodes: List[DAGOrderHelper] = list(filter(lambda n: n.layer_number == layer_idx, total_nodes))
            current_level_nodes.sort(key=lambda n: n.sort_number)

            # 构建当前层节点
            ypos_acc = 0
            current_graphics_nodes = []
            for node in current_level_nodes:
                if node.is_fake_node():
                    curr_gnode = PenetrateNode(20, None)
                    curr_gnode.setPos(previous_node_end, ypos_acc)
                    curr_gnode.attach_data(node)
                    ypos_acc += curr_gnode.boundingRect().height()
                    current_graphics_nodes.append(curr_gnode)
                    self.scene_bind.addItem(curr_gnode)
                else:
                    if node.layer_bind.bind_point().is_start:
                        curr_gnode = ActivePresentNode(node.layer_bind.bind_point().point_name, StoryNodeType.StoryStartNode,
                                                       self.font(), None)
                    else:
                        curr_gnode = ActivePresentNode(node.layer_bind.bind_point().point_name, StoryNodeType.FragmentNode,
                                                       self.font(), None)
                    curr_gnode.attach_data(node)
                    curr_gnode.setPos(previous_node_end, ypos_acc)
                    ypos_acc += curr_gnode.boundingRect().height()
                    current_graphics_nodes.append(curr_gnode)
                    self.scene_bind.addItem(curr_gnode)
                    pass

                ypos_acc += self.node_span
                pass

            # 调整同层节点宽度
            curr_layer_width = 0
            for n in current_graphics_nodes:
                curr_layer_width = max(curr_layer_width, n.boundingRect().width())
                pass
            for n in current_graphics_nodes:
                if hasattr(n, "resize_width"):
                    n.resize_width(curr_layer_width)
                    pass
                pass

            previous_node_end += curr_layer_width + self.layer_span
            if len(previous_graphics_nodes) > 0:
                prev_layer_width = 0
                for n in previous_graphics_nodes:
                    prev_layer_width = max(prev_layer_width, n.boundingRect().width())
                    pass

                for curr_gnode in current_graphics_nodes:
                    sort_helper: DAGOrderHelper = curr_gnode.get_data()
                    for prev_gnode in previous_graphics_nodes:
                        if prev_gnode.get_data() in sort_helper.get_upstream_nodes():
                            line_cmbn = TransitionCurve(prev_gnode, curr_gnode, prev_layer_width,None)
                            self.scene_bind.addItem(line_cmbn)
                            line_cmbn.layout_refresh()
                            pass
                        pass
                    pass
                pass

            previous_graphics_nodes = current_graphics_nodes
            pass

        pass


if __name__ == "__main__":
    app = QApplication(sys.argv)
    view = DAGActiveView(Direction.RankLR, None)
    view.show()

    arrows = [
        Arrow(Point('a'), Point('b')),
        Arrow(Point('a'), Point('c')),
        Arrow(Point('c'), Point('d')),
        Arrow(Point('a'), Point('d')),
        Arrow(Point('c'), Point('e')),
        Arrow(Point('c'), Point('f')),
    ]
    view.update_with_edges(arrows)

    app.exec()