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.8',           # Padding around the graph
                'nodesep': '3.0',       # Increased spacing between nodes horizontally
                'ranksep': '2.5'        # Increased spacing between ranks vertically
            }
        )
        
        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,
                pos="6,2!"  # Centered at top (adjust x-coordinate based on timeline width)
            )
        
        total_events = len(events)
        
        # Create all event nodes first
        for i, event in enumerate(events):
            event_id = event.get('id', f'event_{i}')
            event_label = event.get('label', f'Event {i+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 position in timeline
            if total_events == 1:
                opacity = 'FF'
            else:
                opacity_value = int(255 * (1.0 - (i * 0.7 / (total_events - 1))))
                opacity = format(opacity_value, '02x')
            
            node_color = f"{base_color}{opacity}"
            font_color = 'white' if i < total_events * 0.7 else 'black'
            
            # Calculate position for serpentine layout
            row = i // events_per_row
            col = i % events_per_row
            
            # For odd rows, reverse the column position to create serpentine effect
            if row % 2 == 1:
                visual_col = events_per_row - 1 - col
            else:
                visual_col = col
            
            # Add the event node with position attributes for layout
            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',
                pos=f"{visual_col * 4.5},{-row * 3}!"  # Increased spacing for serpentine layout
            )
        
        # Connect events in chronological order (1→2→3→4...)
        for i in range(len(events) - 1):
            current_event_id = events[i].get('id', f'event_{i}')
            next_event_id = events[i + 1].get('id', f'event_{i + 1}')
            
            dot.edge(
                current_event_id,
                next_event_id,
                color='#666666',
                arrowsize='0.8',
                penwidth='2'
            )
        
        # No need to connect title to events - it stays at the top independently
        
        # Set the layout engine to handle fixed positions
        dot.engine = 'neato'

        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)}"