Mehvar/modules/visualizer.py

import networkx as nx
from pyvis.network import Network
import textwrap
import ast
import re


class Visualizer:
    def __init__(self):
        pass

    def build_graph(self, propositions, relations, scores, output_path):
        """
        دریافت ۳ ساختار مجزا (گزاره‌ها، روابط، امتیازها) و رسم گراف
        """
        net = Network(height="100vh", width="100%", bgcolor="#1a1a1a", font_color="white", directed=True,
                      cdn_resources="in_line")

        # ساخت دیکشنری‌های کمکی برای جستجوی سریع
        prop_dict = {p['proposition_id']: p['proposition_text'] for p in propositions}
        score_dict = {s['relation_id']: s['score'] for s in scores}

        # 1. اضافه کردن نودها (گزاره‌ها)
        for p_id, text in prop_dict.items():
            wrapped_text = "\n".join(textwrap.wrap(text, width=40))
            net.add_node(
                p_id,
                label=f"گزاره {p_id}",
                title=wrapped_text,  # نمایش متن کامل هنگام هاور پیش‌فرض
                shape="box",
                color={"background": "#264653", "border": "#1b3039", "highlight": {"background": "#2a9d8f"}},
                font={"face": "B Nazanin, Tahoma, Arial", "size": 16, "color": "white"}
            )

        # 2. اضافه کردن یال‌ها (روابط)
        for rel in relations:
            source = rel['source_proposition_id']
            target = rel['target_proposition_id']
            rel_id = rel['relation_id']
            rel_type = rel['relation_type']
            score = score_dict.get(rel_id, 0)

            wrapped_rel = "\n".join(textwrap.wrap(rel_type, width=40))
            edge_title = f"امتیاز: {score}\n\nنوع رابطه:\n{wrapped_rel}"

            if source in prop_dict and target in prop_dict:
                net.add_edge(
                    source,
                    target,
                    label=str(score),
                    title=edge_title,  # نمایش متن در سایدبار از طریق رویداد
                    width=max(2, int(score) / 2),  # ضخامت خط بر اساس امتیاز
                    color={"color": "#e9c46a", "highlight": "#f4a261"},
                    font={"align": "middle", "color": "#ffea00", "background": "#1a1a1a"},
                    arrows={"to": {"enabled": True, "scaleFactor": 1.2}}
                )

        # تنظیمات گراف و فعال‌سازی هاور
        net.set_options("""
        {
          "physics": { 
              "forceAtlas2Based": { 
                  "gravitationalConstant": -50, 
                  "centralGravity": 0.01, 
                  "springLength": 150, 
                  "springConstant": 0.08 
              }, 
              "minVelocity": 0.75, 
              "solver": "forceAtlas2Based" 
          },
          "interaction": { "hover": true, "navigationButtons": true }
        }
        """)

        html_content = net.generate_html(output_path)

        # تزریق اسکریپت برای ارسال اطلاعات هاور به سایدبار صفحه وب
        hover_js = """
        <script type="text/javascript">
            network.on("hoverEdge", function (params) {
                var edgeData = edges.get(params.edge);
                window.parent.postMessage({ type: 'hoverEdge', info: edgeData.title }, '*');
            });
            network.on("hoverNode", function (params) {
                var nodeData = nodes.get(params.node);
                window.parent.postMessage({ type: 'hoverNode', info: nodeData.title }, '*');
            });
            network.on("blurEdge", function (params) {
                window.parent.postMessage({ type: 'blur', info: '' }, '*');
            });
            network.on("blurNode", function (params) {
                window.parent.postMessage({ type: 'blur', info: '' }, '*');
            });
        </script>
        </body>
        """
        html_content = html_content.replace('</body>', hover_js)

        with open(output_path, "w", encoding="utf-8") as f:
            f.write(html_content)
        return output_path

    def build_tree(self, tree_data_list, output_path):
        net = Network(height="100vh", width="100%", bgcolor="#1a1a1a", font_color="white", directed=True,
                      cdn_resources="in_line")
        valid_node_ids = set()

        for node_data in tree_data_list:
            node_id = str(node_data.get('segment_id', '')).strip()
            valid_node_ids.add(node_id)
            level = int(node_data.get('level', 0))

            full_text = node_data.get('semantic_unit', 'بدون متن')
            wrapped_text = "\n".join(textwrap.wrap(full_text, width=40))

            raw_sentences = node_data.get('original_sentences', [])
            if isinstance(raw_sentences, str):
                try:
                    raw_sentences = ast.literal_eval(raw_sentences)
                except (ValueError, SyntaxError):
                    pass

            if isinstance(raw_sentences, list):
                sentences_str = "<br><br>".join(raw_sentences)
                num_sentences = len(raw_sentences)
            else:
                sentences_str = str(raw_sentences)
                num_sentences = len([s for s in sentences_str.split('.') if s.strip()]) if sentences_str else 0

            # متمایز کردن رنگ برگ‌ها (سطح 0) از نودهای میانی
            bg_color = "#e76f51" if level == 0 else "#2a9d8f"

            net.add_node(
                node_id,
                label=f"سطح {level}\n{node_id}",
                title=wrapped_text,
                level=level,
                shape="box",
                sentences=sentences_str,
                num_sentences=num_sentences,
                color={"background": bg_color, "border": "#1f6e64", "highlight": {"background": "#3bc2b4"}},
                font={"face": "B Nazanin, Tahoma, Arial", "size": 16}
            )

        for node_data in tree_data_list:
            node_id = str(node_data.get('segment_id', '')).strip()
            parent_val = node_data.get('parent_segment_id')

            if parent_val and str(parent_val).strip() != "":
                clean_parents_str = re.sub(r"[\[\]\'\"\s]", "", str(parent_val))
                parents_list = clean_parents_str.split(',')

                for p_id in parents_list:
                    if p_id in valid_node_ids:
                        net.add_edge(
                            p_id,
                            node_id,
                            width=3,
                            color={"color": "#e9c46a", "highlight": "#f4a261"},
                            arrows={"to": {"enabled": True, "scaleFactor": 1.2}},
                            smooth={'type': 'cubicBezier', 'forceDirection': 'vertical', 'roundness': 0.4}
                        )

        net.set_options("""
        {
          "layout": { 
              "hierarchical": { 
                  "enabled": true, 
                  "direction": "DU", 
                  "sortMethod": "directed", 
                  "nodeSpacing": 250, 
                  "levelSeparation": 150 
              } 
          },
          "physics": { 
              "hierarchicalRepulsion": { 
                  "centralGravity": 0.0, 
                  "springLength": 100, 
                  "springConstant": 0.01, 
                  "nodeDistance": 250, 
                  "damping": 0.09 
              }, 
              "solver": "hierarchicalRepulsion" 
          },
          "interaction": { "hover": true, "navigationButtons": true }
        }
        """)

        html_content = net.generate_html(output_path)

        # تزریق اسکریپت کلیک درخت
        click_js = """
        <script type="text/javascript">
            network.on("click", function (params) {
                if (params.nodes.length > 0) {
                    var nodeId = params.nodes[0];
                    var clickedNode = nodes.get(nodeId);
                    window.parent.postMessage({
                        type: 'node_click',
                        id: nodeId,
                        label: clickedNode.label,
                        summary: clickedNode.title,
                        sentences: clickedNode.sentences || 'اطلاعات جملات موجود نیست.',
                        num_sentences: clickedNode.num_sentences || 0
                    }, '*');
                }
            });
        </script>
        </body>
        """
        html_content = html_content.replace('</body>', click_js)

        with open(output_path, "w", encoding="utf-8") as f:
            f.write(html_content)

        return output_path