app.py

import gradio as gr
import networkx as nx
from pyvis.network import Network
import base64
import tempfile
import os

# Entity type colours (matching your NER tool)
# Updated for better distinction between Person, Location, and Event
ENTITY_COLOURS = {
    'PERSON': '#00B894',      # Green
    'LOCATION': '#9B59B6',    # Purple (more distinct from green)
    'EVENT': '#F39C12',       # Orange/Gold (very distinct)
    'ORGANIZATION': '#55A3FF', # Light Blue
    'DATE': '#FF6B6B'         # Red
}

# Relationship types for dropdown (alphabetical, capitalised, no underscores)
RELATIONSHIP_TYPES = [
    'Attended',
    'Born in',
    'Collaborates with',
    'Died in',
    'Employed by',
    'Founded',
    'Knows',
    'Lives at',
    'Located in',
    'Married to',
    'Member of',
    'Occurred at',
    'Other',
    'Parent of',
    'Participated in',
    'Related to',
    'Sibling of',
    'Visited',
    'Works with',
    'Wrote'
]


class NetworkGraphBuilder:
    def __init__(self):
        self.entities = []
        self.relationships = []
        
    def add_entity(self, name, entity_type, record_id):
        """Add an entity to the collection"""
        if name and name.strip():
            # Avoid duplicates
            existing = [e for e in self.entities if e['name'].lower() == name.strip().lower()]
            if not existing:
                self.entities.append({
                    'name': name.strip(),
                    'type': entity_type,
                    'record_id': record_id
                })
    
    def add_relationship(self, source, target, rel_type):
        """Add a relationship between entities"""
        if source and target and source.strip() and target.strip() and source.strip() != target.strip():
            self.relationships.append({
                'source': source.strip(),
                'target': target.strip(),
                'type': rel_type if rel_type else 'Related to'
            })
    
    def build_graph(self):
        """Build NetworkX graph from entities and relationships"""
        G = nx.Graph()
        
        # Add nodes with attributes
        for entity in self.entities:
            G.add_node(
                entity['name'],
                entity_type=entity['type'],
                record_id=entity['record_id']
            )
        
        # Add edges
        for rel in self.relationships:
            if rel['source'] in G.nodes and rel['target'] in G.nodes:
                G.add_edge(
                    rel['source'],
                    rel['target'],
                    relationship=rel['type']
                )
        
        return G
    
    def create_pyvis_graph(self, G, for_export=False):
        """Create interactive PyVis visualisation. Returns (iframe_html, standalone_html, stats_dict)"""
        if len(G.nodes) == 0:
            return None, None, None
        
        # Potato SVG as base64 data URI (Easter egg for Excellent Boiled Potatoes!)
        potato_svg = '''<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 100 100" width="50" height="50">
            <ellipse cx="50" cy="50" rx="40" ry="30" fill="#C4A484" transform="rotate(-15 50 50)"/>
            <ellipse cx="35" cy="40" rx="4" ry="3" fill="#8B7355"/>
            <ellipse cx="55" cy="35" rx="3" ry="4" fill="#8B7355"/>
            <ellipse cx="65" cy="50" rx="4" ry="3" fill="#8B7355"/>
            <ellipse cx="45" cy="60" rx="3" ry="4" fill="#8B7355"/>
            <ellipse cx="30" cy="55" rx="3" ry="2" fill="#8B7355"/>
        </svg>'''
        potato_data_uri = "data:image/svg+xml;base64," + base64.b64encode(potato_svg.encode()).decode()
        
        # Set height based on whether this is for export
        graph_height = "800px" if for_export else "600px"
        
        # Create PyVis network with dark theme
        net = Network(
            height=graph_height,
            width="100%",
            bgcolor="#1a1a2e",
            font_color="white",
            directed=False
        )
        
        # Configure physics and styling
        net.set_options("""
        {
            "nodes": {
                "borderWidth": 2,
                "borderWidthSelected": 4,
                "font": {
                    "size": 18,
                    "face": "arial",
                    "color": "white",
                    "strokeWidth": 3,
                    "strokeColor": "#1a1a2e"
                },
                "shadow": {
                    "enabled": true,
                    "color": "rgba(0,0,0,0.5)",
                    "size": 10
                }
            },
            "edges": {
                "color": {
                    "color": "#888888",
                    "highlight": "#ffffff"
                },
                "font": {
                    "size": 14,
                    "face": "arial",
                    "color": "#ffffff",
                    "strokeWidth": 0,
                    "background": "rgba(26, 26, 46, 0.8)",
                    "align": "middle"
                },
                "smooth": {
                    "enabled": true,
                    "type": "continuous"
                },
                "width": 2
            },
            "physics": {
                "enabled": true,
                "barnesHut": {
                    "gravitationalConstant": -8000,
                    "centralGravity": 0.3,
                    "springLength": 150,
                    "springConstant": 0.04,
                    "damping": 0.09
                },
                "stabilisation": {
                    "enabled": true,
                    "iterations": 100
                }
            },
            "interaction": {
                "hover": true,
                "tooltipDelay": 200,
                "dragNodes": true,
                "dragView": true,
                "zoomView": true
            }
        }
        """)
        
        # Add nodes with styling based on entity type and degree
        for node in G.nodes():
            data = G.nodes[node]
            entity_type = data.get('entity_type', 'UNKNOWN')
            colour = ENTITY_COLOURS.get(entity_type, '#CCCCCC')
            
            # Size based on connections (degree)
            degree = G.degree(node)
            size = 25 + (degree * 8)
            
            # Create plain text tooltip (no HTML)
            connections = list(G.neighbors(node))
            title_lines = [
                node,
                f"Type: {entity_type}",
                f"Connections: {len(connections)}"
            ]
            if connections:
                connected_to = ', '.join(connections[:5])
                if len(connections) > 5:
                    connected_to += f"... +{len(connections) - 5} more"
                title_lines.append(f"Connected to: {connected_to}")
            
            title = '\n'.join(title_lines)
            
            # Check if this is a potato node (Easter egg!)
            is_potato = 'potato' in node.lower()
            
            if is_potato:
                net.add_node(
                    node,
                    label=node,
                    shape="image",
                    image=potato_data_uri,
                    size=size + 10,
                    title=title,
                    font={'size': 18, 'color': 'white', 'strokeWidth': 3, 'strokeColor': '#1a1a2e'}
                )
            else:
                net.add_node(
                    node,
                    label=node,
                    color=colour,
                    size=size,
                    title=title,
                    font={'size': 18, 'color': 'white', 'strokeWidth': 3, 'strokeColor': '#1a1a2e'}
                )
        
        # Add edges with relationship labels
        for edge in G.edges(data=True):
            rel_type = edge[2].get('relationship', '')
            net.add_edge(
                edge[0],
                edge[1],
                title=rel_type,
                label=rel_type,
                color='#888888',
                width=2,
                font={'size': 14, 'color': '#ffffff', 'strokeWidth': 0, 'background': 'rgba(26,26,46,0.8)'}
            )
        
        # Generate HTML
        html = net.generate_html()
        
        # Store standalone HTML for export
        standalone_html = html
        
        # Calculate stats for export
        stats_dict = {
            'nodes': G.number_of_nodes(),
            'edges': G.number_of_edges(),
            'density': nx.density(G) if len(G.edges) > 0 else 0,
            'avg_degree': sum(dict(G.degree()).values()) / G.number_of_nodes() if G.number_of_nodes() > 0 else 0,
            'top_nodes': sorted(dict(G.degree()).items(), key=lambda x: x[1], reverse=True)[:3] if G.number_of_nodes() > 0 else []
        }
        
        # Encode as base64 data URI for iframe src
        html_bytes = html.encode('utf-8')
        b64_html = base64.b64encode(html_bytes).decode('utf-8')
        
        iframe_html = f'''
        <iframe 
            src="data:text/html;base64,{b64_html}" 
            width="100%" 
            height="620px" 
            style="border: 2px solid #333; border-radius: 10px; background-color: #1a1a2e;"
            sandbox="allow-scripts allow-same-origin"
        ></iframe>
        '''
        
        return iframe_html, standalone_html, stats_dict


def collect_entities_from_records(
    p1, l1, e1, o1, d1,
    p2, l2, e2, o2, d2,
    p3, l3, e3, o3, d3,
    p4, l4, e4, o4, d4,
    p5, l5, e5, o5, d5,
    p6, l6, e6, o6, d6,
    p7, l7, e7, o7, d7,
    p8, l8, e8, o8, d8,
    p9, l9, e9, o9, d9,
    p10, l10, e10, o10, d10
):
    """Collect all entities from the input fields"""
    builder = NetworkGraphBuilder()
    
    # Process each record
    records = [
        (p1, l1, e1, o1, d1),
        (p2, l2, e2, o2, d2),
        (p3, l3, e3, o3, d3),
        (p4, l4, e4, o4, d4),
        (p5, l5, e5, o5, d5),
        (p6, l6, e6, o6, d6),
        (p7, l7, e7, o7, d7),
        (p8, l8, e8, o8, d8),
        (p9, l9, e9, o9, d9),
        (p10, l10, e10, o10, d10),
    ]
    
    for record_id, (person, location, event, org, date) in enumerate(records, 1):
        if person:
            builder.add_entity(person, 'PERSON', record_id)
        if location:
            builder.add_entity(location, 'LOCATION', record_id)
        if event:
            builder.add_entity(event, 'EVENT', record_id)
        if org:
            builder.add_entity(org, 'ORGANIZATION', record_id)
        if date:
            builder.add_entity(date, 'DATE', record_id)
    
    # Create list of all entity names for relationship dropdowns
    entity_names = sorted([e['name'] for e in builder.entities])
    
    # Count by type
    counts = {
        'PERSON': sum(1 for e in builder.entities if e['type'] == 'PERSON'),
        'LOCATION': sum(1 for e in builder.entities if e['type'] == 'LOCATION'),
        'EVENT': sum(1 for e in builder.entities if e['type'] == 'EVENT'),
        'ORGANIZATION': sum(1 for e in builder.entities if e['type'] == 'ORGANIZATION'),
        'DATE': sum(1 for e in builder.entities if e['type'] == 'DATE'),
    }
    
    # Create HTML summary that spans full width
    summary_html = f'''
    <div style="background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); padding: 20px; border-radius: 12px; margin: 10px 0;">
        <h3 style="color: white; margin: 0 0 15px 0; font-size: 18px;">📊 Identified Entities ({len(builder.entities)} total)</h3>
        <div style="display: flex; flex-wrap: wrap; gap: 15px; margin-bottom: 15px;">
            <div style="background: rgba(255,255,255,0.2); padding: 10px 20px; border-radius: 8px; text-align: center;">
                <div style="font-size: 24px; font-weight: bold; color: white;">{counts['PERSON']}</div>
                <div style="color: rgba(255,255,255,0.9); font-size: 12px;">👤 People</div>
            </div>
            <div style="background: rgba(255,255,255,0.2); padding: 10px 20px; border-radius: 8px; text-align: center;">
                <div style="font-size: 24px; font-weight: bold; color: white;">{counts['LOCATION']}</div>
                <div style="color: rgba(255,255,255,0.9); font-size: 12px;">📍 Locations</div>
            </div>
            <div style="background: rgba(255,255,255,0.2); padding: 10px 20px; border-radius: 8px; text-align: center;">
                <div style="font-size: 24px; font-weight: bold; color: white;">{counts['EVENT']}</div>
                <div style="color: rgba(255,255,255,0.9); font-size: 12px;">📅 Events</div>
            </div>
            <div style="background: rgba(255,255,255,0.2); padding: 10px 20px; border-radius: 8px; text-align: center;">
                <div style="font-size: 24px; font-weight: bold; color: white;">{counts['ORGANIZATION']}</div>
                <div style="color: rgba(255,255,255,0.9); font-size: 12px;">🏢 Organisations</div>
            </div>
            <div style="background: rgba(255,255,255,0.2); padding: 10px 20px; border-radius: 8px; text-align: center;">
                <div style="font-size: 24px; font-weight: bold; color: white;">{counts['DATE']}</div>
                <div style="color: rgba(255,255,255,0.9); font-size: 12px;">🗓️ Dates</div>
            </div>
        </div>
        <div style="background: rgba(255,255,255,0.15); padding: 12px; border-radius: 8px;">
            <div style="color: rgba(255,255,255,0.9); font-size: 13px; word-wrap: break-word;">
                <strong>Entities:</strong> {', '.join(entity_names) if entity_names else 'None found'}
            </div>
        </div>
    </div>
    '''
    
    # Return summary and update all 20 dropdowns (10 source + 10 target)
    return (
        summary_html,
        gr.update(choices=entity_names, value=None),  # source 1
        gr.update(choices=entity_names, value=None),  # target 1
        gr.update(choices=entity_names, value=None),  # source 2
        gr.update(choices=entity_names, value=None),  # target 2
        gr.update(choices=entity_names, value=None),  # source 3
        gr.update(choices=entity_names, value=None),  # target 3
        gr.update(choices=entity_names, value=None),  # source 4
        gr.update(choices=entity_names, value=None),  # target 4
        gr.update(choices=entity_names, value=None),  # source 5
        gr.update(choices=entity_names, value=None),  # target 5
        gr.update(choices=entity_names, value=None),  # source 6
        gr.update(choices=entity_names, value=None),  # target 6
        gr.update(choices=entity_names, value=None),  # source 7
        gr.update(choices=entity_names, value=None),  # target 7
        gr.update(choices=entity_names, value=None),  # source 8
        gr.update(choices=entity_names, value=None),  # target 8
        gr.update(choices=entity_names, value=None),  # source 9
        gr.update(choices=entity_names, value=None),  # target 9
        gr.update(choices=entity_names, value=None),  # source 10
        gr.update(choices=entity_names, value=None),  # target 10
    )


def generate_network_graph(
    p1, l1, e1, o1, d1,
    p2, l2, e2, o2, d2,
    p3, l3, e3, o3, d3,
    p4, l4, e4, o4, d4,
    p5, l5, e5, o5, d5,
    p6, l6, e6, o6, d6,
    p7, l7, e7, o7, d7,
    p8, l8, e8, o8, d8,
    p9, l9, e9, o9, d9,
    p10, l10, e10, o10, d10,
    src1, rel1, tgt1,
    src2, rel2, tgt2,
    src3, rel3, tgt3,
    src4, rel4, tgt4,
    src5, rel5, tgt5,
    src6, rel6, tgt6,
    src7, rel7, tgt7,
    src8, rel8, tgt8,
    src9, rel9, tgt9,
    src10, rel10, tgt10
):
    """Generate the network graph from all inputs"""
    try:
        builder = NetworkGraphBuilder()
        
        # Process each record
        records = [
            (p1, l1, e1, o1, d1),
            (p2, l2, e2, o2, d2),
            (p3, l3, e3, o3, d3),
            (p4, l4, e4, o4, d4),
            (p5, l5, e5, o5, d5),
            (p6, l6, e6, o6, d6),
            (p7, l7, e7, o7, d7),
            (p8, l8, e8, o8, d8),
            (p9, l9, e9, o9, d9),
            (p10, l10, e10, o10, d10),
        ]
        
        for record_id, (person, location, event, org, date) in enumerate(records, 1):
            if person:
                builder.add_entity(person, 'PERSON', record_id)
            if location:
                builder.add_entity(location, 'LOCATION', record_id)
            if event:
                builder.add_entity(event, 'EVENT', record_id)
            if org:
                builder.add_entity(org, 'ORGANIZATION', record_id)
            if date:
                builder.add_entity(date, 'DATE', record_id)
        
        # Process relationships (now 10 total)
        relationships = [
            (src1, rel1, tgt1),
            (src2, rel2, tgt2),
            (src3, rel3, tgt3),
            (src4, rel4, tgt4),
            (src5, rel5, tgt5),
            (src6, rel6, tgt6),
            (src7, rel7, tgt7),
            (src8, rel8, tgt8),
            (src9, rel9, tgt9),
            (src10, rel10, tgt10),
        ]
        
        for source, rel_type, target in relationships:
            if source and target:
                builder.add_relationship(source, target, rel_type)
        
        # Build graph
        G = builder.build_graph()
        
        if len(G.nodes) == 0:
            empty_html = '''
            <div style="background-color: #1a1a2e; height: 400px; display: flex; align-items: center; justify-content: center; border-radius: 10px; border: 2px solid #333;">
                <div style="text-align: center; color: white;">
                    <div style="font-size: 48px; margin-bottom: 20px;">📊</div>
                    <h3>No entities to display</h3>
                    <p style="color: #888;">Enter entities above and click "Process Entities" first</p>
                </div>
            </div>
            '''
            return empty_html, "❌ **No entities to display.** Please enter entities in Step 1 first."
        
        # Create visualisation
        graph_html, standalone_html, stats_dict = builder.create_pyvis_graph(G)
        
        # Create statistics with colour legend included
        stats_html = f'''
        <div style="background: #f8f9fa; padding: 15px; border-radius: 10px; border: 1px solid #ddd;">
            <h4 style="margin: 0 0 10px 0;">📈 Network Statistics</h4>
            <table style="width: 100%; border-collapse: collapse;">
                <tr>
                    <td style="padding: 8px; border-bottom: 1px solid #eee;"><strong>Nodes (Entities)</strong></td>
                    <td style="padding: 8px; border-bottom: 1px solid #eee; text-align: right;">{G.number_of_nodes()}</td>
                </tr>
                <tr>
                    <td style="padding: 8px; border-bottom: 1px solid #eee;"><strong>Edges (Relationships)</strong></td>
                    <td style="padding: 8px; border-bottom: 1px solid #eee; text-align: right;">{G.number_of_edges()}</td>
                </tr>
        '''
        
        if len(G.edges) > 0:
            density = nx.density(G)
            avg_degree = sum(dict(G.degree()).values()) / G.number_of_nodes()
            stats_html += f'''
                <tr>
                    <td style="padding: 8px; border-bottom: 1px solid #eee;">
                        <strong>Density (Connectedness)</strong>
                        <div style="font-size: 11px; color: #888; font-weight: normal;">0 = isolated, 1 = fully linked</div>
                    </td>
                    <td style="padding: 8px; border-bottom: 1px solid #eee; text-align: right;">{density:.3f}</td>
                </tr>
                <tr>
                    <td style="padding: 8px;"><strong>Avg. Density</strong></td>
                    <td style="padding: 8px; text-align: right;">{avg_degree:.2f}</td>
                </tr>
            </table>
            '''
            
            # Most connected
            degrees = dict(G.degree())
            top_nodes = sorted(degrees.items(), key=lambda x: x[1], reverse=True)[:3]
            stats_html += '''
            <div style="margin-top: 15px;">
                <strong>Most Connected:</strong>
                <ul style="margin: 5px 0; padding-left: 20px;">
            '''
            for node, degree in top_nodes:
                stats_html += f'<li>{node}: {degree}</li>'
            stats_html += '</ul></div>'
        else:
            stats_html += '''
                </table>
                <div style="margin-top: 10px; padding: 10px; background: #fff3cd; border-radius: 5px; color: #856404;">
                    ⚠️ No relationships defined - nodes are isolated
                </div>
            '''
        
        stats_html += '</div>'
        
        # Add colour legend below stats
        stats_html += f'''
        <div style="background: linear-gradient(135deg, #1a1a2e 0%, #16213e 100%); padding: 15px; border-radius: 10px; margin-top: 15px;">
            <h4 style="color: white; margin: 0 0 12px 0; font-size: 14px;">🎨 Entity Colour Legend</h4>
            <div style="display: flex; flex-direction: column; gap: 8px;">
                <span style="display: flex; align-items: center; gap: 8px; color: white; font-size: 13px;">
                    <span style="width: 16px; height: 16px; border-radius: 50%; background-color: {ENTITY_COLOURS['PERSON']}; display: inline-block; border: 2px solid white;"></span>
                    Person
                </span>
                <span style="display: flex; align-items: center; gap: 8px; color: white; font-size: 13px;">
                    <span style="width: 16px; height: 16px; border-radius: 50%; background-color: {ENTITY_COLOURS['LOCATION']}; display: inline-block; border: 2px solid white;"></span>
                    Location
                </span>
                <span style="display: flex; align-items: center; gap: 8px; color: white; font-size: 13px;">
                    <span style="width: 16px; height: 16px; border-radius: 50%; background-color: {ENTITY_COLOURS['EVENT']}; display: inline-block; border: 2px solid white;"></span>
                    Event
                </span>
                <span style="display: flex; align-items: center; gap: 8px; color: white; font-size: 13px;">
                    <span style="width: 16px; height: 16px; border-radius: 50%; background-color: {ENTITY_COLOURS['ORGANIZATION']}; display: inline-block; border: 2px solid white;"></span>
                    Organisation
                </span>
                <span style="display: flex; align-items: center; gap: 8px; color: white; font-size: 13px;">
                    <span style="width: 16px; height: 16px; border-radius: 50%; background-color: {ENTITY_COLOURS['DATE']}; display: inline-block; border: 2px solid white;"></span>
                    Date
                </span>
            </div>
            <p style="color: white; margin: 12px 0 0 0; font-size: 12px;">
                🖱️ <strong style="color: white;">Interaction:</strong> Drag nodes • Scroll to zoom • Hover for details
            </p>
        </div>
        '''
        
        return graph_html, stats_html
    
    except Exception as e:
        import traceback
        error_trace = traceback.format_exc()
        
        error_html = f'''
        <div style="background-color: #f8d7da; height: 200px; display: flex; align-items: center; justify-content: center; border-radius: 10px; border: 2px solid #f5c6cb;">
            <div style="text-align: center; color: #721c24; padding: 20px;">
                <div style="font-size: 48px; margin-bottom: 10px;">⚠️</div>
                <h3>Error generating graph</h3>
                <p>{str(e)}</p>
            </div>
        </div>
        '''
        return error_html, f"❌ Error: {str(e)}"


def export_network_graph(
    p1, l1, e1, o1, d1,
    p2, l2, e2, o2, d2,
    p3, l3, e3, o3, d3,
    p4, l4, e4, o4, d4,
    p5, l5, e5, o5, d5,
    p6, l6, e6, o6, d6,
    p7, l7, e7, o7, d7,
    p8, l8, e8, o8, d8,
    p9, l9, e9, o9, d9,
    p10, l10, e10, o10, d10,
    src1, rel1, tgt1,
    src2, rel2, tgt2,
    src3, rel3, tgt3,
    src4, rel4, tgt4,
    src5, rel5, tgt5,
    src6, rel6, tgt6,
    src7, rel7, tgt7,
    src8, rel8, tgt8,
    src9, rel9, tgt9,
    src10, rel10, tgt10
):
    """Export the network graph as a standalone HTML file"""
    try:
        builder = NetworkGraphBuilder()
        
        # Process each record
        records = [
            (p1, l1, e1, o1, d1),
            (p2, l2, e2, o2, d2),
            (p3, l3, e3, o3, d3),
            (p4, l4, e4, o4, d4),
            (p5, l5, e5, o5, d5),
            (p6, l6, e6, o6, d6),
            (p7, l7, e7, o7, d7),
            (p8, l8, e8, o8, d8),
            (p9, l9, e9, o9, d9),
            (p10, l10, e10, o10, d10),
        ]
        
        for record_id, (person, location, event, org, date) in enumerate(records, 1):
            if person:
                builder.add_entity(person, 'PERSON', record_id)
            if location:
                builder.add_entity(location, 'LOCATION', record_id)
            if event:
                builder.add_entity(event, 'EVENT', record_id)
            if org:
                builder.add_entity(org, 'ORGANIZATION', record_id)
            if date:
                builder.add_entity(date, 'DATE', record_id)
        
        # Process relationships
        relationships = [
            (src1, rel1, tgt1),
            (src2, rel2, tgt2),
            (src3, rel3, tgt3),
            (src4, rel4, tgt4),
            (src5, rel5, tgt5),
            (src6, rel6, tgt6),
            (src7, rel7, tgt7),
            (src8, rel8, tgt8),
            (src9, rel9, tgt9),
            (src10, rel10, tgt10),
        ]
        
        for source, rel_type, target in relationships:
            if source and target:
                builder.add_relationship(source, target, rel_type)
        
        # Build graph
        G = builder.build_graph()
        
        if len(G.nodes) == 0:
            return None
        
        # Create visualisation with export flag for larger size
        _, standalone_html, stats_dict = builder.create_pyvis_graph(G, for_export=True)
        
        # Create enhanced HTML with legend and statistics
        if standalone_html:
            # Build statistics HTML
            stats_content = f'''
            <div style="margin-bottom: 20px;">
                <h3 style="margin: 0 0 15px 0; color: #333;">📈 Network Statistics</h3>
                <table style="width: 100%; border-collapse: collapse; background: white; border-radius: 8px; overflow: hidden; box-shadow: 0 2px 4px rgba(0,0,0,0.1);">
                    <tr style="background: #f8f9fa;">
                        <td style="padding: 12px; border-bottom: 1px solid #eee;"><strong>Nodes (Entities)</strong></td>
                        <td style="padding: 12px; border-bottom: 1px solid #eee; text-align: right;">{stats_dict['nodes']}</td>
                    </tr>
                    <tr>
                        <td style="padding: 12px; border-bottom: 1px solid #eee;"><strong>Edges (Relationships)</strong></td>
                        <td style="padding: 12px; border-bottom: 1px solid #eee; text-align: right;">{stats_dict['edges']}</td>
                    </tr>
                    <tr style="background: #f8f9fa;">
                        <td style="padding: 12px; border-bottom: 1px solid #eee;"><strong>Density</strong></td>
                        <td style="padding: 12px; border-bottom: 1px solid #eee; text-align: right;">{stats_dict['density']:.3f}</td>
                    </tr>
                    <tr>
                        <td style="padding: 12px;"><strong>Avg. Connections</strong></td>
                        <td style="padding: 12px; text-align: right;">{stats_dict['avg_degree']:.2f}</td>
                    </tr>
                </table>
            '''
            
            if stats_dict['top_nodes']:
                stats_content += '''
                <div style="margin-top: 15px; background: white; padding: 15px; border-radius: 8px; box-shadow: 0 2px 4px rgba(0,0,0,0.1);">
                    <strong>Most Connected:</strong>
                    <ul style="margin: 10px 0 0 0; padding-left: 20px;">
                '''
                for node, degree in stats_dict['top_nodes']:
                    stats_content += f'<li>{node}: {degree} connections</li>'
                stats_content += '</ul></div>'
            
            stats_content += '</div>'
            
            # Build legend HTML
            legend_html = f'''
            <div style="background: linear-gradient(135deg, #1a1a2e 0%, #16213e 100%); padding: 20px; border-radius: 10px; margin-bottom: 20px;">
                <h3 style="color: white; margin: 0 0 15px 0; font-size: 16px;">🎨 Entity Colour Legend</h3>
                <div style="display: flex; flex-wrap: wrap; gap: 20px;">
                    <span style="display: flex; align-items: center; gap: 8px; color: white;">
                        <span style="width: 20px; height: 20px; border-radius: 50%; background-color: {ENTITY_COLOURS['PERSON']}; display: inline-block; border: 2px solid white;"></span>
                        Person
                    </span>
                    <span style="display: flex; align-items: center; gap: 8px; color: white;">
                        <span style="width: 20px; height: 20px; border-radius: 50%; background-color: {ENTITY_COLOURS['LOCATION']}; display: inline-block; border: 2px solid white;"></span>
                        Location
                    </span>
                    <span style="display: flex; align-items: center; gap: 8px; color: white;">
                        <span style="width: 20px; height: 20px; border-radius: 50%; background-color: {ENTITY_COLOURS['EVENT']}; display: inline-block; border: 2px solid white;"></span>
                        Event
                    </span>
                    <span style="display: flex; align-items: center; gap: 8px; color: white;">
                        <span style="width: 20px; height: 20px; border-radius: 50%; background-color: {ENTITY_COLOURS['ORGANIZATION']}; display: inline-block; border: 2px solid white;"></span>
                        Organisation
                    </span>
                    <span style="display: flex; align-items: center; gap: 8px; color: white;">
                        <span style="width: 20px; height: 20px; border-radius: 50%; background-color: {ENTITY_COLOURS['DATE']}; display: inline-block; border: 2px solid white;"></span>
                        Date
                    </span>
                </div>
                <p style="color: white; margin: 15px 0 0 0; font-size: 13px;">
                    🖱️ <strong style="color: white;">Interaction:</strong> Drag nodes • Scroll to zoom • Hover for details
                </p>
            </div>
            '''
            
            # Encode the graph HTML as base64 for the iframe
            graph_b64 = base64.b64encode(standalone_html.encode('utf-8')).decode('utf-8')
            
            # Create enhanced HTML wrapper
            enhanced_html = f'''<!DOCTYPE html>
<html>
<head>
    <meta charset="utf-8">
    <title>Network Graph - Basic Network Explorer</title>
    <style>
        body {{
            font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, Oxygen, Ubuntu, sans-serif;
            margin: 0;
            padding: 20px;
            background-color: #f5f5f5;
        }}
        .container {{
            max-width: 1600px;
            margin: 0 auto;
        }}
        .header {{
            text-align: center;
            margin-bottom: 20px;
            padding: 20px;
            background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
            border-radius: 12px;
            color: white;
        }}
        .header h1 {{
            margin: 0 0 10px 0;
        }}
        .header p {{
            margin: 0;
            opacity: 0.9;
        }}
        .content {{
            display: flex;
            gap: 20px;
        }}
        .graph-container {{
            flex: 3;
            background: white;
            border-radius: 12px;
            padding: 15px;
            box-shadow: 0 4px 6px rgba(0,0,0,0.1);
        }}
        .sidebar {{
            flex: 1;
            min-width: 300px;
        }}
        .graph-frame {{
            width: 100%;
            height: 800px;
            border: 2px solid #333;
            border-radius: 8px;
            background-color: #1a1a2e;
        }}
        .footer {{
            text-align: center;
            margin-top: 20px;
            padding: 15px;
            color: #666;
            font-size: 13px;
        }}
        @media (max-width: 1000px) {{
            .content {{
                flex-direction: column;
            }}
            .sidebar {{
                min-width: 100%;
            }}
        }}
    </style>
</head>
<body>
    <div class="container">
        <div class="header">
            <h1>🕸️ Network Graph</h1>
            <p>Created with Basic Network Explorer</p>
        </div>
        <div class="content">
            <div class="graph-container">
                <iframe class="graph-frame" src="data:text/html;base64,{graph_b64}"></iframe>
            </div>
            <div class="sidebar">
                {stats_content}
                {legend_html}
            </div>
        </div>
        <div class="footer">
            <p>Exported from <strong>Basic Network Explorer</strong> — Bodleian Libraries (Oxford) Sassoon Research Fellowship</p>
        </div>
    </div>
</body>
</html>'''
            
            export_file = tempfile.NamedTemporaryFile(
                mode='w',
                suffix='.html',
                prefix='network_graph_',
                delete=False,
                encoding='utf-8'
            )
            export_file.write(enhanced_html)
            export_file.close()
            return export_file.name
        
        return None
    
    except Exception as e:
        return None


def load_austen_example():
    """Load the Jane Austen Pride and Prejudice example"""
    return (
        # Record 1
        "Elizabeth Bennet", "Longbourn", "Meryton Ball", "", "",
        # Record 2
        "Mr. Darcy", "Pemberley", "Netherfield Ball", "", "",
        # Record 3
        "Jane Bennet", "Netherfield", "", "", "",
        # Record 4
        "Mr. Bingley", "London", "", "", "",
        # Record 5
        "Mr. Wickham", "", "", "Militia", "",
        # Record 6
        "Charlotte Lucas", "", "", "", "",
        # Record 7
        "Mr. Collins", "", "Excellent Boiled Potatoes", "", "",
        # Record 8
        "Caroline Bingley", "", "", "", "",
        # Record 9
        "Lydia Bennet", "", "", "", "",
        # Record 10
        "Lady Catherine de Bourgh", "", "", "", "",
    )


def load_wwii_example():
    """Load a WWII history example"""
    return (
        # Record 1
        "Winston Churchill", "London", "Battle of Britain", "War Cabinet", "1940",
        # Record 2
        "Franklin D. Roosevelt", "Washington D.C.", "D-Day", "Allied Forces", "1944",
        # Record 3
        "Field Marshal Montgomery", "North Africa", "Battle of El Alamein", "Eighth Army", "1942",
        # Record 4
        "Clement Attlee", "Potsdam", "Potsdam Conference", "Labour Party", "1945",
        # Record 5
        "", "", "", "", "",
        # Record 6
        "", "", "", "", "",
        # Record 7
        "", "", "", "", "",
        # Record 8
        "", "", "", "", "",
        # Record 9
        "", "", "", "", "",
        # Record 10
        "", "", "", "", "",
    )


def create_coloured_label(text, colour, emoji):
    """Create a coloured pill-style label HTML"""
    return f'''
    <span style="
        background-color: {colour}; 
        color: white; 
        padding: 4px 12px; 
        border-radius: 15px; 
        font-size: 13px; 
        font-weight: 500;
        display: inline-block;
        margin-bottom: 5px;
    ">{emoji} {text}</span>
    '''


def create_interface():
    # Note: theme moved to launch() for Gradio 6.0 compatibility
    with gr.Blocks(title="Network Explorer") as demo:
        gr.Markdown("""
        # 🕸️ Basic Network Explorer
        
        Build interactive network graphs by entering entities extracted through Named Entity Recognition (NER).
        
        ### How to use this tool:
        1. **📝 Enter entities** in the records below (or load an example to get started)
        2. **⚙️ Click "Process Entities"** to collect and prepare all entities for relationships
        3. **🤝 Define relationships** between entities using the dropdowns (or type your own relationship type)
        4. **🎨 Click "Generate Network Graph"** to visualise
        5. **👁️ Explore** - drag nodes to rearrange, scroll to zoom, hover for details
        6. **💾 Export (optional)** - click "Export as HTML" to download your graph as an interactive file
        """)
        
        gr.HTML("""
        <div style="background-color: #fff3cd; border: 1px solid #ffeaa7; border-radius: 8px; padding: 12px; margin: 15px 0;">
            <strong style="color: #856404;">💡 Top tip:</strong> Start with just a few entities and relationships to see how it works!
        </div>
        """)
        
        # Quick start buttons
        gr.Markdown("### 💡 Quick Start - Load an Example:")
        with gr.Row():
            austen_btn = gr.Button("📚 Jane Austen (Pride & Prejudice)", variant="secondary", size="sm")
            wwii_btn = gr.Button("⚔️ WWII History", variant="secondary", size="sm")
        
        gr.HTML("<hr style='margin: 15px 0;'>")
        
        # ==================== STEP 1: ENTITY INPUT (5 per row) ====================
        gr.Markdown("## 📝 Step 1: Enter Entities")
        
        entity_inputs = []
        
        # First row: Records 1-5
        with gr.Row():
            with gr.Column(scale=1, min_width=180):
                gr.Markdown("**Record 1**")
                gr.HTML(create_coloured_label("Person", ENTITY_COLOURS['PERSON'], "👤"))
                p1 = gr.Textbox(label="", placeholder="e.g. Elizabeth Bennet", show_label=False)
                gr.HTML(create_coloured_label("Location", ENTITY_COLOURS['LOCATION'], "📍"))
                l1 = gr.Textbox(label="", placeholder="e.g. Longbourn", show_label=False)
                gr.HTML(create_coloured_label("Event", ENTITY_COLOURS['EVENT'], "📅"))
                e1 = gr.Textbox(label="", placeholder="e.g. Meryton Ball", show_label=False)
                gr.HTML(create_coloured_label("Organisation", ENTITY_COLOURS['ORGANIZATION'], "🏢"))
                o1 = gr.Textbox(label="", placeholder="e.g. The Militia", show_label=False)
                gr.HTML(create_coloured_label("Date", ENTITY_COLOURS['DATE'], "🗓️"))
                d1 = gr.Textbox(label="", placeholder="e.g. 1812", show_label=False)
                entity_inputs.extend([p1, l1, e1, o1, d1])
            
            with gr.Column(scale=1, min_width=180):
                gr.Markdown("**Record 2**")
                gr.HTML(create_coloured_label("Person", ENTITY_COLOURS['PERSON'], "👤"))
                p2 = gr.Textbox(label="", placeholder="e.g. Mr. Darcy", show_label=False)
                gr.HTML(create_coloured_label("Location", ENTITY_COLOURS['LOCATION'], "📍"))
                l2 = gr.Textbox(label="", placeholder="e.g. Pemberley", show_label=False)
                gr.HTML(create_coloured_label("Event", ENTITY_COLOURS['EVENT'], "📅"))
                e2 = gr.Textbox(label="", placeholder="", show_label=False)
                gr.HTML(create_coloured_label("Organisation", ENTITY_COLOURS['ORGANIZATION'], "🏢"))
                o2 = gr.Textbox(label="", placeholder="", show_label=False)
                gr.HTML(create_coloured_label("Date", ENTITY_COLOURS['DATE'], "🗓️"))
                d2 = gr.Textbox(label="", placeholder="", show_label=False)
                entity_inputs.extend([p2, l2, e2, o2, d2])
            
            with gr.Column(scale=1, min_width=180):
                gr.Markdown("**Record 3**")
                gr.HTML(create_coloured_label("Person", ENTITY_COLOURS['PERSON'], "👤"))
                p3 = gr.Textbox(label="", placeholder="", show_label=False)
                gr.HTML(create_coloured_label("Location", ENTITY_COLOURS['LOCATION'], "📍"))
                l3 = gr.Textbox(label="", placeholder="", show_label=False)
                gr.HTML(create_coloured_label("Event", ENTITY_COLOURS['EVENT'], "📅"))
                e3 = gr.Textbox(label="", placeholder="", show_label=False)
                gr.HTML(create_coloured_label("Organisation", ENTITY_COLOURS['ORGANIZATION'], "🏢"))
                o3 = gr.Textbox(label="", placeholder="", show_label=False)
                gr.HTML(create_coloured_label("Date", ENTITY_COLOURS['DATE'], "🗓️"))
                d3 = gr.Textbox(label="", placeholder="", show_label=False)
                entity_inputs.extend([p3, l3, e3, o3, d3])
            
            with gr.Column(scale=1, min_width=180):
                gr.Markdown("**Record 4**")
                gr.HTML(create_coloured_label("Person", ENTITY_COLOURS['PERSON'], "👤"))
                p4 = gr.Textbox(label="", placeholder="", show_label=False)
                gr.HTML(create_coloured_label("Location", ENTITY_COLOURS['LOCATION'], "📍"))
                l4 = gr.Textbox(label="", placeholder="", show_label=False)
                gr.HTML(create_coloured_label("Event", ENTITY_COLOURS['EVENT'], "📅"))
                e4 = gr.Textbox(label="", placeholder="", show_label=False)
                gr.HTML(create_coloured_label("Organisation", ENTITY_COLOURS['ORGANIZATION'], "🏢"))
                o4 = gr.Textbox(label="", placeholder="", show_label=False)
                gr.HTML(create_coloured_label("Date", ENTITY_COLOURS['DATE'], "🗓️"))
                d4 = gr.Textbox(label="", placeholder="", show_label=False)
                entity_inputs.extend([p4, l4, e4, o4, d4])
            
            with gr.Column(scale=1, min_width=180):
                gr.Markdown("**Record 5**")
                gr.HTML(create_coloured_label("Person", ENTITY_COLOURS['PERSON'], "👤"))
                p5 = gr.Textbox(label="", placeholder="", show_label=False)
                gr.HTML(create_coloured_label("Location", ENTITY_COLOURS['LOCATION'], "📍"))
                l5 = gr.Textbox(label="", placeholder="", show_label=False)
                gr.HTML(create_coloured_label("Event", ENTITY_COLOURS['EVENT'], "📅"))
                e5 = gr.Textbox(label="", placeholder="", show_label=False)
                gr.HTML(create_coloured_label("Organisation", ENTITY_COLOURS['ORGANIZATION'], "🏢"))
                o5 = gr.Textbox(label="", placeholder="", show_label=False)
                gr.HTML(create_coloured_label("Date", ENTITY_COLOURS['DATE'], "🗓️"))
                d5 = gr.Textbox(label="", placeholder="", show_label=False)
                entity_inputs.extend([p5, l5, e5, o5, d5])
        
        # Additional records 6-10
        with gr.Accordion("➕ Additional Records (6-10)", open=False):
            with gr.Row():
                with gr.Column(scale=1, min_width=180):
                    gr.Markdown("**Record 6**")
                    gr.HTML(create_coloured_label("Person", ENTITY_COLOURS['PERSON'], "👤"))
                    p6 = gr.Textbox(label="", show_label=False)
                    gr.HTML(create_coloured_label("Location", ENTITY_COLOURS['LOCATION'], "📍"))
                    l6 = gr.Textbox(label="", show_label=False)
                    gr.HTML(create_coloured_label("Event", ENTITY_COLOURS['EVENT'], "📅"))
                    e6 = gr.Textbox(label="", show_label=False)
                    gr.HTML(create_coloured_label("Organisation", ENTITY_COLOURS['ORGANIZATION'], "🏢"))
                    o6 = gr.Textbox(label="", show_label=False)
                    gr.HTML(create_coloured_label("Date", ENTITY_COLOURS['DATE'], "🗓️"))
                    d6 = gr.Textbox(label="", show_label=False)
                    entity_inputs.extend([p6, l6, e6, o6, d6])
                
                with gr.Column(scale=1, min_width=180):
                    gr.Markdown("**Record 7**")
                    gr.HTML(create_coloured_label("Person", ENTITY_COLOURS['PERSON'], "👤"))
                    p7 = gr.Textbox(label="", show_label=False)
                    gr.HTML(create_coloured_label("Location", ENTITY_COLOURS['LOCATION'], "📍"))
                    l7 = gr.Textbox(label="", show_label=False)
                    gr.HTML(create_coloured_label("Event", ENTITY_COLOURS['EVENT'], "📅"))
                    e7 = gr.Textbox(label="", show_label=False)
                    gr.HTML(create_coloured_label("Organisation", ENTITY_COLOURS['ORGANIZATION'], "🏢"))
                    o7 = gr.Textbox(label="", show_label=False)
                    gr.HTML(create_coloured_label("Date", ENTITY_COLOURS['DATE'], "🗓️"))
                    d7 = gr.Textbox(label="", show_label=False)
                    entity_inputs.extend([p7, l7, e7, o7, d7])
                
                with gr.Column(scale=1, min_width=180):
                    gr.Markdown("**Record 8**")
                    gr.HTML(create_coloured_label("Person", ENTITY_COLOURS['PERSON'], "👤"))
                    p8 = gr.Textbox(label="", show_label=False)
                    gr.HTML(create_coloured_label("Location", ENTITY_COLOURS['LOCATION'], "📍"))
                    l8 = gr.Textbox(label="", show_label=False)
                    gr.HTML(create_coloured_label("Event", ENTITY_COLOURS['EVENT'], "📅"))
                    e8 = gr.Textbox(label="", show_label=False)
                    gr.HTML(create_coloured_label("Organisation", ENTITY_COLOURS['ORGANIZATION'], "🏢"))
                    o8 = gr.Textbox(label="", show_label=False)
                    gr.HTML(create_coloured_label("Date", ENTITY_COLOURS['DATE'], "🗓️"))
                    d8 = gr.Textbox(label="", show_label=False)
                    entity_inputs.extend([p8, l8, e8, o8, d8])
                
                with gr.Column(scale=1, min_width=180):
                    gr.Markdown("**Record 9**")
                    gr.HTML(create_coloured_label("Person", ENTITY_COLOURS['PERSON'], "👤"))
                    p9 = gr.Textbox(label="", show_label=False)
                    gr.HTML(create_coloured_label("Location", ENTITY_COLOURS['LOCATION'], "📍"))
                    l9 = gr.Textbox(label="", show_label=False)
                    gr.HTML(create_coloured_label("Event", ENTITY_COLOURS['EVENT'], "📅"))
                    e9 = gr.Textbox(label="", show_label=False)
                    gr.HTML(create_coloured_label("Organisation", ENTITY_COLOURS['ORGANIZATION'], "🏢"))
                    o9 = gr.Textbox(label="", show_label=False)
                    gr.HTML(create_coloured_label("Date", ENTITY_COLOURS['DATE'], "🗓️"))
                    d9 = gr.Textbox(label="", show_label=False)
                    entity_inputs.extend([p9, l9, e9, o9, d9])
                
                with gr.Column(scale=1, min_width=180):
                    gr.Markdown("**Record 10**")
                    gr.HTML(create_coloured_label("Person", ENTITY_COLOURS['PERSON'], "👤"))
                    p10 = gr.Textbox(label="", show_label=False)
                    gr.HTML(create_coloured_label("Location", ENTITY_COLOURS['LOCATION'], "📍"))
                    l10 = gr.Textbox(label="", show_label=False)
                    gr.HTML(create_coloured_label("Event", ENTITY_COLOURS['EVENT'], "📅"))
                    e10 = gr.Textbox(label="", show_label=False)
                    gr.HTML(create_coloured_label("Organisation", ENTITY_COLOURS['ORGANIZATION'], "🏢"))
                    o10 = gr.Textbox(label="", show_label=False)
                    gr.HTML(create_coloured_label("Date", ENTITY_COLOURS['DATE'], "🗓️"))
                    d10 = gr.Textbox(label="", show_label=False)
                    entity_inputs.extend([p10, l10, e10, o10, d10])
        
        gr.HTML("<hr style='margin: 20px 0;'>")
        
        # ==================== STEP 2: PROCESS ENTITIES ====================
        gr.Markdown("## ⚙️ Step 2: Process Entities")
        gr.Markdown("*Click the button below to collect all entities before defining relationships*")
        
        # Process button
        collect_btn = gr.Button("⚙️ Process Entities", variant="primary", size="lg")
        
        # Full-width entity summary
        entity_summary = gr.HTML()
        
        gr.HTML("<hr style='margin: 20px 0;'>")
        
        # ==================== STEP 3: RELATIONSHIPS (Grid) ====================
        gr.Markdown("## 🤝 Step 3: Define Relationships")
        gr.Markdown("*Select entities from the dropdowns to create connections (click 'Process Entities' first). You can also type your own relationship type.*")
        
        # Relationship inputs in a grid (5 columns for first 5)
        relationship_inputs = []
        
        with gr.Row():
            with gr.Column(scale=1, min_width=180):
                gr.Markdown("**Relationship 1**")
                src1 = gr.Dropdown(label="From", choices=[])
                rel1 = gr.Dropdown(label="Type", choices=RELATIONSHIP_TYPES, value="Related to", allow_custom_value=True)
                tgt1 = gr.Dropdown(label="To", choices=[])
                relationship_inputs.extend([src1, rel1, tgt1])
            
            with gr.Column(scale=1, min_width=180):
                gr.Markdown("**Relationship 2**")
                src2 = gr.Dropdown(label="From", choices=[])
                rel2 = gr.Dropdown(label="Type", choices=RELATIONSHIP_TYPES, value="Related to", allow_custom_value=True)
                tgt2 = gr.Dropdown(label="To", choices=[])
                relationship_inputs.extend([src2, rel2, tgt2])
            
            with gr.Column(scale=1, min_width=180):
                gr.Markdown("**Relationship 3**")
                src3 = gr.Dropdown(label="From", choices=[])
                rel3 = gr.Dropdown(label="Type", choices=RELATIONSHIP_TYPES, value="Related to", allow_custom_value=True)
                tgt3 = gr.Dropdown(label="To", choices=[])
                relationship_inputs.extend([src3, rel3, tgt3])
            
            with gr.Column(scale=1, min_width=180):
                gr.Markdown("**Relationship 4**")
                src4 = gr.Dropdown(label="From", choices=[])
                rel4 = gr.Dropdown(label="Type", choices=RELATIONSHIP_TYPES, value="Related to", allow_custom_value=True)
                tgt4 = gr.Dropdown(label="To", choices=[])
                relationship_inputs.extend([src4, rel4, tgt4])
            
            with gr.Column(scale=1, min_width=180):
                gr.Markdown("**Relationship 5**")
                src5 = gr.Dropdown(label="From", choices=[])
                rel5 = gr.Dropdown(label="Type", choices=RELATIONSHIP_TYPES, value="Related to", allow_custom_value=True)
                tgt5 = gr.Dropdown(label="To", choices=[])
                relationship_inputs.extend([src5, rel5, tgt5])
        
        # Additional relationships 6-10
        with gr.Accordion("➕ Additional Relationships (6-10)", open=False):
            with gr.Row():
                with gr.Column(scale=1, min_width=180):
                    gr.Markdown("**Relationship 6**")
                    src6 = gr.Dropdown(label="From", choices=[])
                    rel6 = gr.Dropdown(label="Type", choices=RELATIONSHIP_TYPES, value="Related to", allow_custom_value=True)
                    tgt6 = gr.Dropdown(label="To", choices=[])
                    relationship_inputs.extend([src6, rel6, tgt6])
                
                with gr.Column(scale=1, min_width=180):
                    gr.Markdown("**Relationship 7**")
                    src7 = gr.Dropdown(label="From", choices=[])
                    rel7 = gr.Dropdown(label="Type", choices=RELATIONSHIP_TYPES, value="Related to", allow_custom_value=True)
                    tgt7 = gr.Dropdown(label="To", choices=[])
                    relationship_inputs.extend([src7, rel7, tgt7])
                
                with gr.Column(scale=1, min_width=180):
                    gr.Markdown("**Relationship 8**")
                    src8 = gr.Dropdown(label="From", choices=[])
                    rel8 = gr.Dropdown(label="Type", choices=RELATIONSHIP_TYPES, value="Related to", allow_custom_value=True)
                    tgt8 = gr.Dropdown(label="To", choices=[])
                    relationship_inputs.extend([src8, rel8, tgt8])
                
                with gr.Column(scale=1, min_width=180):
                    gr.Markdown("**Relationship 9**")
                    src9 = gr.Dropdown(label="From", choices=[])
                    rel9 = gr.Dropdown(label="Type", choices=RELATIONSHIP_TYPES, value="Related to", allow_custom_value=True)
                    tgt9 = gr.Dropdown(label="To", choices=[])
                    relationship_inputs.extend([src9, rel9, tgt9])
                
                with gr.Column(scale=1, min_width=180):
                    gr.Markdown("**Relationship 10**")
                    src10 = gr.Dropdown(label="From", choices=[])
                    rel10 = gr.Dropdown(label="Type", choices=RELATIONSHIP_TYPES, value="Related to", allow_custom_value=True)
                    tgt10 = gr.Dropdown(label="To", choices=[])
                    relationship_inputs.extend([src10, rel10, tgt10])
        
        gr.HTML("<hr style='margin: 20px 0;'>")
        
        # ==================== STEP 4: GENERATE & VIEW ====================
        gr.Markdown("## 🎨 Step 4: Generate Network Graph")
        
        generate_btn = gr.Button("🎨 Generate Network Graph", variant="primary", size="lg")
        
        # Full-width network graph with stats sidebar (legend now included in stats)
        with gr.Row():
            with gr.Column(scale=3):
                network_plot = gr.HTML(label="Interactive Network Graph")
            with gr.Column(scale=1):
                network_stats = gr.HTML()
        
        # Export section
        gr.Markdown("### 💾 Export Your Graph")
        gr.Markdown("*Generate a graph first, then click export to download*")
        with gr.Row():
            export_btn = gr.Button("💾 Export as HTML", variant="secondary", size="sm")
            export_file = gr.File(
                label="Download Interactive HTML",
                file_types=[".html"],
                interactive=False
            )
        gr.HTML("""
        <div style="background-color: #e8f4f8; border: 1px solid #bee5eb; border-radius: 8px; padding: 12px; margin: 10px 0;">
            <span style="color: #0c5460;">💡 The exported HTML file is fully interactive — open it in any web browser to explore your network!</span>
        </div>
        """)
        
        # ==================== WIRE UP EVENTS ====================
        
        # Example buttons
        austen_btn.click(
            fn=load_austen_example,
            inputs=[],
            outputs=entity_inputs
        )
        
        wwii_btn.click(
            fn=load_wwii_example,
            inputs=[],
            outputs=entity_inputs
        )
        
        # Collect entities - now updates 20 dropdowns (10 relationships x 2)
        collect_btn.click(
            fn=collect_entities_from_records,
            inputs=entity_inputs,
            outputs=[
                entity_summary,
                src1, tgt1,
                src2, tgt2,
                src3, tgt3,
                src4, tgt4,
                src5, tgt5,
                src6, tgt6,
                src7, tgt7,
                src8, tgt8,
                src9, tgt9,
                src10, tgt10
            ]
        )
        
        # Generate graph - outputs the visualisation and stats
        all_inputs = entity_inputs + relationship_inputs
        generate_btn.click(
            fn=generate_network_graph,
            inputs=all_inputs,
            outputs=[network_plot, network_stats]
        )
        
        # Export graph - separate button for downloading
        export_btn.click(
            fn=export_network_graph,
            inputs=all_inputs,
            outputs=[export_file]
        )
        
        # Model Information & Documentation section
        gr.HTML("""
        <hr style="margin: 40px 0 20px 0;">
        <div style="background-color: #f8f9fa; border: 1px solid #e9ecef; border-radius: 8px; padding: 20px; margin: 20px 0;">
            <h3 style="margin: 0 0 10px 0;">📚 Library Information & Documentation</h3>
            <p style="color: #666; margin-bottom: 15px;">Learn more about the libraries used in this tool:</p>
            <ul style="list-style-type: circle; padding-left: 20px; margin: 0;">
                <li style="margin-bottom: 8px;">
                    <strong>NetworkX:</strong> 
                    <a href="https://networkx.org/documentation/stable/" target="_blank" style="color: #667eea;">NetworkX Documentation ↗</a>
                    <span style="color: #888; font-size: 13px;"> — Python library for creating, manipulating, and studying complex networks</span>
                </li>
                <li style="margin-bottom: 8px;">
                    <strong>PyVis:</strong> 
                    <a href="https://pyvis.readthedocs.io/en/latest/" target="_blank" style="color: #667eea;">PyVis Documentation ↗</a>
                    <span style="color: #888; font-size: 13px;"> — Interactive network visualisation library built on vis.js</span>
                </li>
                <li style="margin-bottom: 8px;">
                    <strong>Gradio:</strong> 
                    <a href="https://www.gradio.app/docs" target="_blank" style="color: #667eea;">Gradio Documentation ↗</a>
                    <span style="color: #888; font-size: 13px;"> — Web interface framework for machine learning demos</span>
                </li>
            </ul>
        </div>
        """)
        
        # Footer
        gr.HTML("""
        <hr style="margin: 20px 0 20px 0;">
        <div style="text-align: center; color: #666; font-size: 14px; padding: 20px;">
            <p>This <strong>Basic Network Explorer</strong> tool was developed as part of a Bodleian Libraries (Oxford) Sassoon Research Fellowship.</p>
            <p style="color: #888; font-size: 12px; margin-top: 10px;">Built with the aid of Claude Opus 4.5.</p>
        </div>
        """)
    
    return demo


if __name__ == "__main__":
    demo = create_interface()
    # Gradio 6.0 compatibility: theme and ssr_mode moved to launch()
    demo.launch(ssr_mode=False, theme=gr.themes.Soft())