src/utils/mindmap_generator.py

"""
Mindmap Generator using PyVis for interactive radial visualizations
Creates beautiful, explorable mindmaps with custom styling
"""

from pyvis.network import Network
import networkx as nx
from typing import Dict, Any, Optional
import streamlit.components.v1 as components
import tempfile
import os


class MindmapGenerator:
    """
    Generate interactive radial mindmaps using PyVis

    Features:
    - Radial layout with ForceAtlas2 physics
    - Color-coded hierarchy levels
    - Interactive zoom, pan, and hover
    - Customizable styling and dimensions
    """

    def __init__(
        self,
        height: str = "650px",
        width: str = "100%",
        bgcolor: str = "#1e1e1e",
        font_color: str = "#ffffff"
    ):
        """
        Initialize mindmap generator with display settings

        Args:
            height: Height of visualization (CSS format)
            width: Width of visualization (CSS format)
            bgcolor: Background color (hex)
            font_color: Font color (hex)
        """
        self.height = height
        self.width = width
        self.bgcolor = bgcolor
        self.font_color = font_color

        # Color scheme for node levels
        self.level_colors = {
            0: "#ff6b6b",  # Center - red/coral
            1: "#4ecdc4",  # Primary - teal
            2: "#95e1d3",  # Secondary - light teal
            3: "#f9ca24",  # Tertiary - yellow
            4: "#a29bfe"   # Quaternary - purple
        }

    def create_radial_mindmap(self, mindmap_data: Dict[str, Any]) -> Network:
        """
        Create interactive radial mindmap from structured data

        Args:
            mindmap_data: Dictionary with center, nodes, and edges

        Returns:
            Configured PyVis Network object
        """
        # Initialize PyVis network
        net = Network(
            height=self.height,
            width=self.width,
            bgcolor=self.bgcolor,
            font_color=self.font_color,
            notebook=False,
            directed=False
        )

        # Configure physics for radial layout
        net.set_options("""
        {
            "physics": {
                "enabled": true,
                "forceAtlas2Based": {
                    "gravitationalConstant": -50,
                    "centralGravity": 0.005,
                    "springLength": 150,
                    "springConstant": 0.08,
                    "damping": 0.4,
                    "avoidOverlap": 0.5
                },
                "maxVelocity": 50,
                "solver": "forceAtlas2Based",
                "timestep": 0.35,
                "stabilization": {
                    "enabled": true,
                    "iterations": 1000,
                    "updateInterval": 25
                }
            },
            "interaction": {
                "hover": true,
                "hoverConnectedEdges": true,
                "tooltipDelay": 100,
                "navigationButtons": true,
                "keyboard": {
                    "enabled": true,
                    "speed": {"x": 10, "y": 10, "zoom": 0.02}
                },
                "zoomView": true,
                "dragView": true
            },
            "edges": {
                "smooth": {
                    "enabled": true,
                    "type": "continuous",
                    "roundness": 0.5
                }
            }
        }
        """)

        # Add center node
        center = mindmap_data.get('center', 'Unknown Topic')
        net.add_node(
            'center',
            label=center,
            title=f"<div style='padding:10px'><b style='font-size:16px'>{center}</b><br><i>Central Topic</i></div>",
            size=45,
            color=self.level_colors[0],
            font={'size': 24, 'color': self.font_color, 'face': 'Arial', 'bold': True},
            borderWidth=3,
            borderWidthSelected=5
        )

        # Add nodes with level-based styling
        for node in mindmap_data.get('nodes', []):
            node_id = node.get('id', '')
            label = node.get('label', node_id)
            level = node.get('level', 1)
            description = node.get('description', '')

            # Size decreases with level
            size = max(35 - (level * 7), 15)

            # Font size decreases with level
            font_size = max(18 - (level * 2), 12)

            # Get color for this level
            color = self.level_colors.get(level, self.level_colors[2])

            # Create rich tooltip
            tooltip = f"""
            <div style='padding:12px; max-width:300px;'>
                <b style='font-size:14px; color:#4ecdc4;'>{label}</b><br>
                <span style='color:#888;'>Level {level}</span><br>
                <p style='margin-top:8px; font-size:12px;'>{description}</p>
            </div>
            """

            net.add_node(
                node_id,
                label=label,
                title=tooltip,
                size=size,
                color=color,
                font={
                    'size': font_size,
                    'color': self.font_color,
                    'face': 'Arial'
                },
                borderWidth=2,
                borderWidthSelected=4,
                shape='dot'
            )

        # Add edges with relationship labels
        for edge in mindmap_data.get('edges', []):
            from_node = edge.get('from', '')
            to_node = edge.get('to', '')
            label = edge.get('label', '')

            # Edge styling
            net.add_edge(
                from_node,
                to_node,
                title=label if label else 'related to',
                label=label if len(label) < 15 else '',  # Only show short labels
                color={'color': '#666666', 'opacity': 0.6},
                width=2,
                smooth={'type': 'continuous'}
            )

        return net

    def generate_html(self, mindmap_data: Dict[str, Any]) -> str:
        """
        Generate HTML string for the mindmap

        Args:
            mindmap_data: Mindmap structure

        Returns:
            Complete HTML as string
        """
        net = self.create_radial_mindmap(mindmap_data)

        # Generate HTML
        html_content = net.generate_html()

        return html_content

    def save_to_file(self, mindmap_data: Dict[str, Any], filename: str = "mindmap.html"):
        """
        Save mindmap to HTML file

        Args:
            mindmap_data: Mindmap structure
            filename: Output filename
        """
        net = self.create_radial_mindmap(mindmap_data)
        net.save_graph(filename)
        print(f"✅ Mindmap saved to {filename}")

    def render_in_streamlit(self, mindmap_data: Dict[str, Any]):
        """
        Render mindmap directly in Streamlit application

        Args:
            mindmap_data: Mindmap structure
        """
        html_content = self.generate_html(mindmap_data)

        # Display using Streamlit components
        components.html(
            html_content,
            height=int(self.height.replace('px', '')),
            scrolling=False
        )


# Convenience functions for direct usage
def generate_mindmap_html(mindmap_data: Dict[str, Any]) -> str:
    """
    Quick function to generate mindmap HTML

    Args:
        mindmap_data: Mindmap structure

    Returns:
        HTML string
    """
    generator = MindmapGenerator()
    return generator.generate_html(mindmap_data)


def create_sample_mindmap() -> Dict[str, Any]:
    """
    Create a sample mindmap for testing

    Returns:
        Sample mindmap data structure
    """
    return {
        'center': 'Machine Learning',
        'nodes': [
            {
                'id': 'supervised',
                'label': 'Supervised Learning',
                'level': 1,
                'description': 'Learning with labeled data'
            },
            {
                'id': 'unsupervised',
                'label': 'Unsupervised Learning',
                'level': 1,
                'description': 'Learning from unlabeled data'
            },
            {
                'id': 'classification',
                'label': 'Classification',
                'level': 2,
                'description': 'Categorizing data into classes'
            },
            {
                'id': 'regression',
                'label': 'Regression',
                'level': 2,
                'description': 'Predicting continuous values'
            }
        ],
        'edges': [
            {'from': 'center', 'to': 'supervised', 'label': 'includes'},
            {'from': 'center', 'to': 'unsupervised', 'label': 'includes'},
            {'from': 'supervised', 'to': 'classification', 'label': 'type'},
            {'from': 'supervised', 'to': 'regression', 'label': 'type'}
        ]
    }


if __name__ == "__main__":
    # Test the generator
    print("Mindmap Generator Module - Ready for import")
    print("Use MindmapGenerator class to create visualizations")

    # Create sample
    sample = create_sample_mindmap()
    generator = MindmapGenerator()
    generator.save_to_file(sample, "test_mindmap.html")