app.py

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

# 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
RELATIONSHIP_TYPES = [
    'related_to',
    'knows',
    'works_with',
    'located_in',
    'participated_in',
    'member_of',
    'occurred_at',
    'employed_by',
    'founded',
    'attended',
    'collaborates_with',
    'married_to',
    'sibling_of',
    'parent_of',
    'lives_at',
    'wrote',
    'visited',
    'born_in',
    'died_in',
    'other'
]


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"""
        if len(G.nodes) == 0:
            return 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()
        
        # 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


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 10 dropdowns (5 source + 5 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
    )


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
):
    """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
        relationships = [
            (src1, rel1, tgt1),
            (src2, rel2, tgt2),
            (src3, rel3, tgt3),
            (src4, rel4, tgt4),
            (src5, rel5, tgt5),
        ]
        
        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 "Identify Entities" first</p>
                </div>
            </div>
            '''
            return empty_html, "❌ **No entities to display.** Please enter entities in Step 1 first."
        
        # Create visualisation
        graph_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</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</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</strong></td>
                    <td style="padding: 8px; border-bottom: 1px solid #eee; text-align: right;">{density:.3f}</td>
                </tr>
                <tr>
                    <td style="padding: 8px;"><strong>Avg. Connections</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 load_austen_example():
    """Load the Jane Austen Pride and Prejudice example"""
    return (
        # Record 1
        "Elizabeth Bennet", "Longbourn", "Meryton Ball", "Bennet Family", "1811",
        # Record 2
        "Mr. Darcy", "Pemberley", "Netherfield Ball", "Darcy Estate", "",
        # Record 3
        "Jane Bennet", "Netherfield", "", "", "",
        # Record 4
        "Mr. Bingley", "London", "", "", "",
        # Record 5
        "Mr. Wickham", "Meryton", "", "Militia", "",
        # Record 6
        "Charlotte Lucas", "Hunsford", "", "", "",
        # Record 7
        "", "", "", "", "",
        # Record 8
        "", "", "", "", "",
    )


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 "Identify Entities"** to collect and list all entities
        3. **🤝 Define relationships** between entities using the dropdowns
        4. **🎨 Click "Generate Network Graph"** to visualise
        5. **👁️ Explore** - drag nodes to rearrange, scroll to zoom, hover for details
        """)
        
        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])
        
        # Identify button
        collect_btn = gr.Button("🔍 Identify Entities", variant="primary", size="lg")
        
        # Full-width entity summary
        entity_summary = gr.HTML()
        
        gr.HTML("<hr style='margin: 20px 0;'>")
        
        # ==================== STEP 2: RELATIONSHIPS (Grid) ====================
        gr.Markdown("## 🤝 Step 2: Define Relationships")
        gr.Markdown("*Select entities from the dropdowns to create connections (click 'Identify Entities' first)*")
        
        # Relationship inputs in a grid (5 columns)
        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")
                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")
                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")
                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")
                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")
                tgt5 = gr.Dropdown(label="To", choices=[])
                relationship_inputs.extend([src5, rel5, tgt5])
        
        gr.HTML("<hr style='margin: 20px 0;'>")
        
        # ==================== STEP 3: GENERATE & VIEW ====================
        gr.Markdown("## 🎨 Step 3: 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()
        
        # 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: #888; 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
        collect_btn.click(
            fn=collect_entities_from_records,
            inputs=entity_inputs,
            outputs=[
                entity_summary,
                src1, tgt1,
                src2, tgt2,
                src3, tgt3,
                src4, tgt4,
                src5, tgt5
            ]
        )
        
        # Generate graph
        all_inputs = entity_inputs + relationship_inputs
        generate_btn.click(
            fn=generate_network_graph,
            inputs=all_inputs,
            outputs=[network_plot, network_stats]
        )
        
        # 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()