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(" ", ""))
    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')
    # 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}"


# md = '''
# Here is a mind map summarizing the topic of combining machine learning (ML) and computational chemistry (CompChem) for predictive insights into chemical systems:

# **I. Introduction**

# * Machine learning (ML) poised to transform chemical sciences
# * Combining ML and CompChem for predictive insights

# **II. Computational Chemistry (CompChem)**

# * Computational quantum chemistry (CQChem)
# * Methods for generating data sets (e.g., wavefunction theory, correlated wavefunction methods, density functional theory)
# * Representations of systems (e.g., simple, complex, ambiguous)

# **III. Wavefunction Theory Methods**

# * Nonrelativistic time-independent Schrödinger equation
# * Electronic Schrödinger equation
# * Hartree-Fock (HF) approach
# * Correlated wavefunction methods (e.g., extended Hückel theory, neglect of diatomic differential overlap)

# **IV. Density Functional Theory (DFT)**

# * Kinetic energy (KE-) or orbital-free (OF-) DFT
# * Exchange-correlation functional (EC)
# * Kohn-Sham (KS-) DFT
# * Semiempirical methods (e.g., extended Hückel theory, neglect of diatomic differential overlap)

# **V. Semiempirical Methods**

# * Extended Hückel theory
# * Neglect of diatomic differential overlap
# * Semiempirical bond-order potentials (BOPs)
# * Semiempirical nuclear quantum effects (NQEs)

# **VI. Response Properties**

# * Nuclear forces (e.g., F = -Π)
# * Hessian calculations (e.g., second derivative of energy with respect to nuclear positions)
# * Energy conserving forces (e.g., dipole moments)

# **VII. Applications of ML in CompChem**

# * Predicting molecular and material properties
# * Predicting chemical reactions and processes
# * Predicting materials properties (e.g., conductivity, optical properties)
# * Predicting drug design and development

# **VIII. Future Directions**

# * Developing more accurate ML models for CompChem
# * Improving the transferability of ML models between different systems
# * Using ML to accelerate and improve the discovery of new materials and compounds
# '''
# generate_mindmap_svg(md)