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):
        """Create interactive PyVis visualisation. Returns (iframe_html, standalone_html)"""
        if len(G.nodes) == 0:
            return None, None
        
        # Create PyVis network with dark theme
        net = Network(
            height="600px",
            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)
            
            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
        
        # 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


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
):
    """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),
    ]
    
    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 16 dropdowns (8 source + 8 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
    )


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,
    src1, rel1, tgt1,
    src2, rel2, tgt2,
    src3, rel3, tgt3,
    src4, rel4, tgt4,
    src5, rel5, tgt5,
    src6, rel6, tgt6,
    src7, rel7, tgt7,
    src8, rel8, tgt8
):
    """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),
        ]
        
        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 8 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),
        ]
        
        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 = builder.create_pyvis_graph(G)
        
        # Create statistics
        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>'
        
        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,
    src1, rel1, tgt1,
    src2, rel2, tgt2,
    src3, rel3, tgt3,
    src4, rel4, tgt4,
    src5, rel5, tgt5,
    src6, rel6, tgt6,
    src7, rel7, tgt7,
    src8, rel8, tgt8
):
    """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),
        ]
        
        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),
        ]
        
        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
        _, standalone_html = builder.create_pyvis_graph(G)
        
        # Save standalone HTML to a temporary file for download
        if standalone_html:
            export_file = tempfile.NamedTemporaryFile(
                mode='w',
                suffix='.html',
                prefix='network_graph_',
                delete=False,
                encoding='utf-8'
            )
            export_file.write(standalone_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", "", "", "", "",
    )


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


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():
    with gr.Blocks(title="Network Explorer", theme=gr.themes.Soft()) 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 (4 per row) ====================
        gr.Markdown("## 📝 Step 1: Enter Entities")
        
        entity_inputs = []
        
        # First row: Records 1-4 (all on one line)
        with gr.Row():
            with gr.Column(scale=1, min_width=200):
                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., Bennet Family", show_label=False)
                gr.HTML(create_coloured_label("Date", ENTITY_COLOURS['DATE'], "🗓️"))
                d1 = gr.Textbox(label="", placeholder="e.g., 1811", show_label=False)
                entity_inputs.extend([p1, l1, e1, o1, d1])
            
            with gr.Column(scale=1, min_width=200):
                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="e.g., Netherfield Ball", show_label=False)
                gr.HTML(create_coloured_label("Organisation", ENTITY_COLOURS['ORGANIZATION'], "🏢"))
                o2 = gr.Textbox(label="", placeholder="e.g., Darcy Estate", 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=200):
                gr.Markdown("**Record 3**")
                gr.HTML(create_coloured_label("Person", ENTITY_COLOURS['PERSON'], "👤"))
                p3 = gr.Textbox(label="", placeholder="e.g., Jane Bennet", show_label=False)
                gr.HTML(create_coloured_label("Location", ENTITY_COLOURS['LOCATION'], "📍"))
                l3 = gr.Textbox(label="", placeholder="e.g., Netherfield", 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=200):
                gr.Markdown("**Record 4**")
                gr.HTML(create_coloured_label("Person", ENTITY_COLOURS['PERSON'], "👤"))
                p4 = gr.Textbox(label="", placeholder="e.g., Mr. Bingley", show_label=False)
                gr.HTML(create_coloured_label("Location", ENTITY_COLOURS['LOCATION'], "📍"))
                l4 = gr.Textbox(label="", placeholder="e.g., London", 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])
        
        # Additional records 5-8
        with gr.Accordion("➕ Additional Records (5-8)", open=False):
            with gr.Row():
                with gr.Column(scale=1, min_width=200):
                    gr.Markdown("**Record 5**")
                    gr.HTML(create_coloured_label("Person", ENTITY_COLOURS['PERSON'], "👤"))
                    p5 = gr.Textbox(label="", show_label=False)
                    gr.HTML(create_coloured_label("Location", ENTITY_COLOURS['LOCATION'], "📍"))
                    l5 = gr.Textbox(label="", show_label=False)
                    gr.HTML(create_coloured_label("Event", ENTITY_COLOURS['EVENT'], "📅"))
                    e5 = gr.Textbox(label="", show_label=False)
                    gr.HTML(create_coloured_label("Organisation", ENTITY_COLOURS['ORGANIZATION'], "🏢"))
                    o5 = gr.Textbox(label="", show_label=False)
                    gr.HTML(create_coloured_label("Date", ENTITY_COLOURS['DATE'], "🗓️"))
                    d5 = gr.Textbox(label="", show_label=False)
                    entity_inputs.extend([p5, l5, e5, o5, d5])
                
                with gr.Column(scale=1, min_width=200):
                    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=200):
                    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=200):
                    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])
        
        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-8
        with gr.Accordion("➕ Additional Relationships (6-8)", 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])
        
        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
        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>
        """)
        
        # Colour legend
        gr.HTML(f"""
        <div style="background: linear-gradient(135deg, #1a1a2e 0%, #16213e 100%); padding: 20px; border-radius: 10px; margin-top: 20px;">
            <h4 style="color: white; margin: 0 0 15px 0;">🎨 Entity Colour Legend</h4>
            <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>Interaction:</strong> Drag nodes to rearrange • Scroll to zoom • Hover for details
            </p>
        </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 16 dropdowns (8 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
            ]
        )
        
        # 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()
    demo.launch()