import gradio as gr
import networkx as nx
import matplotlib.pyplot as plt
from io import BytesIO

# Function to build and visualize a knowledge graph
def visualize_knowledge_graph(entities, relationships):
    G = nx.DiGraph()
    
    # Add entities as nodes
    for entity in entities.split(","):
        G.add_node(entity.strip())
    
    # Add relationships as edges
    for relationship in relationships.split(";"):
        try:
            source, target, label = map(str.strip, relationship.split(","))
            G.add_edge(source, target, label=label)
        except ValueError:
            return "Invalid relationship format. Use 'entity1,entity2,label' format for each relationship."

    # Plot the graph
    plt.figure(figsize=(10, 6))
    pos = nx.spring_layout(G)
    nx.draw(G, pos, with_labels=True, node_size=3000, node_color="lightblue", font_size=10, font_weight="bold")
    edge_labels = nx.get_edge_attributes(G, 'label')
    nx.draw_networkx_edge_labels(G, pos, edge_labels=edge_labels, font_color="red")
    buf = BytesIO()
    plt.savefig(buf, format="png")
    plt.close()
    buf.seek(0)
    return buf

# Define the Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("## Knowledge Graph Builder")
    gr.Markdown("Enter entities (comma-separated) and relationships (in the format `entity1,entity2,label` separated by semicolons).")
    
    entities_input = gr.Textbox(label="Entities (Comma-separated)", placeholder="e.g., SCIEKore, AI, Education")
    relationships_input = gr.Textbox(label="Relationships (Semicolon-separated)", placeholder="e.g., SCIEKore,AI,uses; SCIEKore,Education,offers")
    
    output_image = gr.Image(label="Knowledge Graph")
    
    submit_button = gr.Button("Build Graph")
    submit_button.click(
        visualize_knowledge_graph, 
        inputs=[entities_input, relationships_input], 
        outputs=output_image
    )

# Launch the Gradio app
if __name__ == "__main__":
    demo.launch(debug=True)