src/viz.py

import tempfile
import os
from typing import Dict, Any
import pyvis.network as net
from .graph import build_tree_structure

def create_pyvis_network(graph: Dict[str, Any], root_id: int) -> str:
    nodes = graph.get("nodes", {})
    nt = net.Network(
        height="600px",
        width="100%",
        bgcolor="#ffffff",
        font_color="black",
        directed=True
    )
    nt.set_options("""
    var options = {
      "physics": {
        "enabled": true,
        "hierarchicalRepulsion": {
          "centralGravity": 0.3,
          "springLength": 120,
          "springConstant": 0.01,
          "nodeDistance": 200,
          "damping": 0.09
        },
        "solver": "hierarchicalRepulsion"
      },
      "layout": {
        "hierarchical": {
          "enabled": true,
          "direction": "UD",
          "sortMethod": "directed"
        }
      }
    }
    """)
    tree = build_tree_structure(graph, root_id)
    colors = ["#ff6b6b", "#4ecdc4", "#45b7d1", "#96ceb4", "#feca57", "#ff9ff3", "#54a0ff"]
    for node_id, node in nodes.items():
        name = node.get("name", str(node_id))
        year_str = f" ({node.get('year')})" if node.get('year') is not None else ""
        label = f"{name}{year_str}"
        depth = tree.get(node_id, {}).get("depth", 0)
        color = colors[depth % len(colors)]
        nt.add_node(
            node_id,
            label=label,
            title=f"Name: {name}\nYear: {node.get('year', 'N/A')}\nInstitution: {node.get('institution', 'N/A')}",
            color=color,
            size=30 if node_id == root_id else 20
        )
    for node_id, node in nodes.items():
        for advisor_id in node.get("advisors", []):
            if advisor_id in nodes:
                nt.add_edge(advisor_id, node_id)
    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".html")
    nt.save_graph(temp_file.name)
    with open(temp_file.name, 'r') as f:
        html_content = f.read()
    os.unlink(temp_file.name)
    return html_content