import json
import pygraphviz as pgv
from io import BytesIO

def generate_mindmap(data):
    # 創(chuàng)建有向圖
    graph = pgv.AGraph(directed=True, layout="dot", rankdir="LR")
    
    # 添加根節(jié)點
    root_id = data["id"]
    graph.add_node(root_id, 
                  label=f"{data['topic']}\n({data['title']})",
                  shape="ellipse",
                  color="red",
                  style="filled",
                  fillcolor="#FFE4B5")
    
    # 遞歸添加子節(jié)點
    def add_children(parent_id, children):
        for child in children:
            child_id = child["id"]
            graph.add_node(child_id, 
                          label=child["topic"],
                          shape="box",
                          color="blue")
            graph.add_edge(parent_id, child_id)
            
            if "children" in child:  # 支持多級子節(jié)點
                add_children(child_id, child["children"])
    
    add_children(root_id, data.get("children", []))
    
    # 生成圖片二進制數(shù)據(jù)
    graph.layout(prog="dot")
    temp = BytesIO()
    graph.draw(temp, format="png")
    temp.seek(0)
    return temp

# 示例數(shù)據(jù)（替換為你的實際數(shù)據(jù)）
method = [
    "create: 創(chuàng)建新節(jié)點",
    "update: 更新節(jié)點內(nèi)容",
    "delete: 刪除指定節(jié)點",
    "move: 移動節(jié)點位置"
]

data = {
    "id": "root",
    "topic": "jsMind",
    "title": "中心主題",
    "children": [
        {
            "id": line.split(":").strip(),
            "topic": line.split(":")[-1].strip(),
        }
        for line in method
        if len(line.split(":").strip()) > 0
    ]
}

# 生成并保存圖片
image_data = generate_mindmap(data)
with open("mindmap.png", "wb") as f:
    f.write(image_data.getvalue())

print("思維導(dǎo)圖已生成：mindmap.png")

import gradio as gr

def visualize_mindmap(method_text):
    method = [line.strip() for line in method_text.split("\n") if line.strip()]
    
    data = {
        "id": "root",
        "topic": "jsMind",
        "title": "中心主題",
        "children": [
            {
                "id": line.split(":").strip(),
                "topic": line.split(":")[-1].strip(),
            }
            for line in method
            if len(line.split(":").strip()) > 0
        ]
    }
    
    return generate_mindmap(data).getvalue()

iface = gr.Interface(
    fn=visualize_mindmap,
    inputs=gr.Textbox(label="輸入方法（每行格式：id: 描述）", lines=5),
    outputs=gr.Image(label="動態(tài)思維導(dǎo)圖"),
    examples=[
        ["create: 創(chuàng)建新節(jié)點\nupdate: 更新節(jié)點內(nèi)容\ndelete: 刪除指定節(jié)點\nmove: 移動節(jié)點位置"]
    ]
)

iface.launch()

pip install pygraphviz
# Windows需額外安裝Graphviz：
# Mac：brew install graphviz
# Linux：sudo apt-get install graphviz