timeline_generator.py

import graphviz
import json
from tempfile import NamedTemporaryFile
import os

def generate_timeline_diagram(json_input: str, output_format: str) -> str:
    """
    Generates a serpentine timeline diagram from JSON input.

    Args:
        json_input (str): A JSON string describing the timeline structure.
                          It must follow the Expected JSON Format Example below.

    Expected JSON Format Example:
    {
      "title": "AI Development Timeline",
      "events_per_row": 4,
      "events": [
        {
          "id": "event_1",
          "label": "Machine Learning Foundations",
          "date": "1950-1960",
          "description": "Early neural networks and perceptrons"
        },
        {
          "id": "event_2",
          "label": "Expert Systems Era",
          "date": "1970-1980",
          "description": "Rule-based AI systems"
        },
        {
          "id": "event_3",
          "label": "Neural Network Revival",
          "date": "1980-1990",
          "description": "Backpropagation algorithm"
        }
      ]
    }

    Returns:
        str: The filepath to the generated PNG image file.
    """
    try:
        if not json_input.strip():
            return "Error: Empty input"
            
        data = json.loads(json_input)
        
        if 'events' not in data:
            raise ValueError("Missing required field: events")

        dot = graphviz.Digraph(
            name='Timeline',
            format='png',
            graph_attr={
                'rankdir': 'TB',        # Top-to-Bottom for better control
                'splines': 'ortho',     # Straight lines with 90-degree bends
                'bgcolor': 'white',     # White background
                'pad': '0.5',           # Padding around the graph
                'nodesep': '1.5',       # Spacing between nodes
                'ranksep': '1.5'        # Spacing between ranks
            }
        )
        
        base_color = '#19191a' # Hardcoded base color
        
        title = data.get('title', '')
        events = data.get('events', [])
        events_per_row = data.get('events_per_row', 4)  # Default to 4 events per row
        
        if not events:
            raise ValueError("Timeline must contain at least one event")

        # Add title node if provided
        if title:
            dot.node(
                'title',
                title,
                shape='plaintext',
                fontsize='18',
                fontweight='bold',
                fontcolor=base_color
            )
        
        # Calculate positions and create serpentine layout
        total_events = len(events)
        previous_event_id = None
        
        # Create invisible nodes for positioning and rank control
        rows = []
        current_row = []
        
        # Group events into rows
        for i, event in enumerate(events):
            current_row.append(event)
            if len(current_row) == events_per_row or i == total_events - 1:
                rows.append(current_row)
                current_row = []
        
        # Process each row and create serpentine connections
        for row_idx, row in enumerate(rows):
            # Determine if row should be reversed (serpentine pattern)
            is_reversed = row_idx % 2 == 1
            if is_reversed:
                row = row[::-1]  # Reverse the row for serpentine effect
            
            # Create invisible nodes for row positioning
            row_nodes = []
            
            for event_idx, event in enumerate(row):
                original_idx = events.index(event)
                event_id = event.get('id', f'event_{original_idx}')
                event_label = event.get('label', f'Event {original_idx+1}')
                event_date = event.get('date', '')
                event_description = event.get('description', '')
                
                # Create full label with date and description
                if event_date and event_description:
                    full_label = f"{event_date}\\n{event_label}\\n{event_description}"
                elif event_date:
                    full_label = f"{event_date}\\n{event_label}"
                elif event_description:
                    full_label = f"{event_label}\\n{event_description}"
                else:
                    full_label = event_label
                
                # Calculate color opacity based on original position in timeline
                if total_events == 1:
                    opacity = 'FF'
                else:
                    opacity_value = int(255 * (1.0 - (original_idx * 0.7 / (total_events - 1))))
                    opacity = format(opacity_value, '02x')
                
                node_color = f"{base_color}{opacity}"
                font_color = 'white' if original_idx < total_events * 0.7 else 'black'
                
                # Add the event node
                dot.node(
                    event_id,
                    full_label,
                    shape='box',
                    style='filled,rounded',
                    fillcolor=node_color,
                    fontcolor=font_color,
                    fontsize='12',
                    width='2.5',
                    height='1.2'
                )
                
                row_nodes.append(event_id)
            
            # Create horizontal connections within the row
            for i in range(len(row_nodes) - 1):
                dot.edge(
                    row_nodes[i],
                    row_nodes[i + 1],
                    color='#666666',
                    arrowsize='0.8',
                    penwidth='2'
                )
            
            # Connect to previous row (serpentine connection)
            if row_idx > 0:
                # Connect last node of previous row to first node of current row
                prev_row_nodes = getattr(generate_timeline_diagram, 'prev_row_nodes', [])
                if prev_row_nodes:
                    # Connect the end of previous row to start of current row
                    if (row_idx - 1) % 2 == 0:  # Previous row was left-to-right
                        connection_start = prev_row_nodes[-1]  # Last node of previous row
                    else:  # Previous row was right-to-left
                        connection_start = prev_row_nodes[0]   # First node of previous row (which was last visually)
                    
                    if row_idx % 2 == 0:  # Current row is left-to-right
                        connection_end = row_nodes[0]  # First node of current row
                    else:  # Current row is right-to-left
                        connection_end = row_nodes[-1]  # Last node of current row (which will be first visually)
                    
                    dot.edge(
                        connection_start,
                        connection_end,
                        color='#666666',
                        arrowsize='0.8',
                        penwidth='2'
                    )
            
            # Store current row nodes for next iteration
            generate_timeline_diagram.prev_row_nodes = row_nodes
            
            # Connect title to first event if title exists and this is the first row
            if title and row_idx == 0:
                first_event = row_nodes[0] if row_idx % 2 == 0 else row_nodes[-1]
                dot.edge('title', first_event, style='invis')
        
        # Clean up the stored attribute
        if hasattr(generate_timeline_diagram, 'prev_row_nodes'):
            delattr(generate_timeline_diagram, 'prev_row_nodes')

        with NamedTemporaryFile(delete=False, suffix=f'.{output_format}') as tmp:
            dot.render(tmp.name, format=output_format, cleanup=True)
            return f"{tmp.name}.{output_format}"

    except json.JSONDecodeError:
        return "Error: Invalid JSON format"
    except Exception as e:
        return f"Error: {str(e)}"