generate_svg.py

import os
from graphviz import Digraph
# from cairosvg import svg2pdf
import re
import os
import dotenv
import random

dotenv.load_dotenv()
hf_token = os.getenv("HF_TOKEN")

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_pdf(svg_file):
#     pdf_file = svg_file.replace(".svg", ".pdf")
#     svg2pdf(file_obj=open(svg_file, "rb"), write_to=pdf_file)
#     return pdf_file

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(" ", ""))
    output_filename = f"{sanitized_title}_mindmap.svg"
    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')
    svg_content = svg_content.replace("%3", root_title)
    # Save the modified SVG content to a file
    with open(output_filename, 'w') as f:
        f.write(svg_content)
    return output_filename

def generate_mindmap(md_text):
    mindmap_svg = generate_mindmap_svg(md_text)
    # mindmap_pdf = generate_mindmap_pdf(mindmap_svg)
    return mindmap_svg

def upload_svg(mindmap_svg):
    from huggingface_hub import HfApi
    api = HfApi(token=hf_token)
    api.upload_file(
        path_or_fileobj=mindmap_svg,
        path_in_repo=f"SVG/{mindmap_svg}",
        repo_id="raannakasturi/ReXploreData",
        repo_type="dataset",
    )
    if os.path.exists(mindmap_svg):
        os.remove(mindmap_svg)
    return f"https://huggingface.co/datasets/raannakasturi/ReXploreData/raw/main/SVG/{mindmap_svg}"

def main(markdown_text):
    mindmap_svg = generate_mindmap_svg(markdown_text.replace("**", ""))
    url = upload_svg(mindmap_svg)
    print(f"Uploaded SVG to {url}")
    return url