generate_mindmap.py

from graphviz import Digraph
import re
import random

def parse_markdown_to_dict(md_text):
    lines = md_text.strip().splitlines()
    mindmap = {}
    stack = []
    for line in lines:
        heading_match = re.match(r'^(#{1,6})\s+(.*)', line)
        bullet_match = re.match(r'^\s*-\s+(.*)', line)
        if heading_match:
            level = len(heading_match.group(1))
            title = heading_match.group(2).strip()
            node = {'title': title, 'children': []}
            while len(stack) >= level:
                stack.pop()
            if stack:
                stack[-1]['children'].append(node)
            else:
                mindmap = node  
            stack.append(node)
        elif bullet_match and stack:
            stack[-1]['children'].append({'title': bullet_match.group(1), 'children': []})
    return mindmap

generated_colors = set()

def generate_random_color():
    """Generate a random color that hasn't been generated before."""
    while True:
        # Generate a random color in hex format
        color = "#{:02x}{:02x}{:02x}".format(random.randint(128, 255), random.randint(128, 255), random.randint(128, 255))
        # If the color is not in the set, it's unique
        if color not in generated_colors:
            generated_colors.add(color)  # Add the color to the set of generated colors
            return color  # Return the unique color
        else:
            continue  # Try again

def brighten_color(color, factor=0.15):
    """Brighten the color by a certain factor (default 10%)"""
    # Remove the '#' symbol
    color = color.lstrip('#')
    
    # Convert hex to RGB
    r, g, b = [int(color[i:i+2], 16) for i in (0, 2, 4)]
    
    # Increase each component by the factor, but clamp to 255
    r = min(255, int(r * (1 + factor)))
    g = min(255, int(g * (1 + factor)))
    b = min(255, int(b * (1 + factor)))
    
    # Convert back to hex
    return "#{:02x}{:02x}{:02x}".format(r, g, b)

def add_nodes_to_graph(graph, node, parent_id=None, font_size=9, parent_color=None):
    node_id = str(id(node))
    title = node['title']
    if parent_color is None:
        node_color = "#ADD8E6"  # Light Blue for the main heading
        border_color = "#000000"  # Dark Blue border for the main heading
        parent_color = "#ADD8E6"
    elif parent_color == "#ADD8E6":
        node_color = generate_random_color()
        border_color = "#808080"
        parent_color = node_color
    else:
        # Child node and its descendants with the same random color
        node_color = brighten_color(parent_color, factor=0.15)
        border_color = "#808080"
    # Check for markdown links
    url_match = re.search(r'\[(.*?)\]\((.*?)\)', title)
    if url_match:
        prefix_text = title[:url_match.start()].strip()
        display_text = url_match.group(1)
        url = url_match.group(2)
        
        label = f'{prefix_text} {display_text}'
        graph.node(node_id, label=label, shape="box", style="rounded,filled", color=border_color, fontcolor="black", fillcolor=node_color, href=url, tooltip=title, fontsize=str(font_size))
    else:
        graph.node(node_id, title, shape="box", style="rounded,filled", color=border_color, fontcolor="black", fillcolor=node_color, tooltip=title, fontsize=str(font_size))
    
    if parent_id:
        graph.edge(parent_id, node_id)

    # Recurse to children, passing down color for the child and its descendants
    for child in node.get('children', []):
        # Assign a random color to each child node (no inheritance from parent)
        add_nodes_to_graph(graph, child, node_id, font_size=max(8, font_size - 1), parent_color=parent_color)

def generate_mindmap_svg(md_text):
    mindmap_dict = parse_markdown_to_dict(md_text)
    root_title = mindmap_dict.get('title', 'Mindmap')
    sanitized_title = re.sub(r'[^a-zA-Z0-9_\-]', '', root_title.replace(" ", ""))
    if output_filename is None:
        output_filename = sanitized_title
    graph = Digraph(format='svg')
    graph.attr(rankdir='LR', size='10,10!', pad="0.5", margin="0.2", ratio="auto")
    graph.attr('node', fontname="Arial", fontsize="9")
    add_nodes_to_graph(graph, mindmap_dict)
    svg_content = graph.pipe(format='svg').decode('utf-8')
    # Replace %3 with the sanitized filename in the SVG content
    svg_content = svg_content.replace("%3", root_title)
    # Save the modified SVG content to a file
    with open(f'{output_filename}.svg', 'w') as f:
        f.write(svg_content)
    return f"{output_filename}".svg