#Developed by ruslanmv.com
import gradio as gr
import matplotlib.pyplot as plt
import networkx as nx
from owlready2 import *
import difflib
import tempfile
import os
def plot_ontology(search_term, owl_file):
    """
    Loads an ontology OWL file, builds a graph of its classes and subclass relationships,
    performs fuzzy matching to find nodes similar to the search term, and returns a plot
    highlighting the matching nodes (and their immediate neighbors) in pink.

    Args:
        search_term (str): The search term used for fuzzy matching on node names.
        owl_file (file): An uploaded ontology file in OWL format.

    Returns:
        A tuple with:
            - The file path of the saved plot image.
            - A text message with search result details.
    """
    if owl_file is None:
        return None, "Please upload an ontology OWL file."

    # Load the ontology from the provided file.
    try:
        # Append temporary directory to the ontology search path.
        onto_path.append(tempfile.gettempdir())
        # Load the ontology using its file URI.
        onto = get_ontology("file://"+owl_file.name).load()
    except Exception as e:
        return None, f"Error loading ontology: {e}"

    # Build a directed graph: nodes represent ontology classes and edges represent subclass relationships.
    G = nx.DiGraph()
    for cls in onto.classes():
        G.add_node(cls.name, iri=cls.iri)
        for parent in cls.is_a:
            if isinstance(parent, ThingClass):
                G.add_edge(parent.name, cls.name)

    # Define a helper function for fuzzy searching nodes.
    def fuzzy_search(query, cutoff=0.6):
        query = query.lower()
        lower_to_orig = {node.lower(): node for node in G.nodes}
        lower_nodes = list(lower_to_orig.keys())
        matches = difflib.get_close_matches(query, lower_nodes, cutoff=cutoff)
        return [lower_to_orig[m] for m in matches]

    search_results = fuzzy_search(search_term)
    if not search_results:
        result_message = f"No matches found for '{search_term}'."
    else:
        result_message = f"Found nodes: {', '.join(search_results)}"

    # Determine which nodes to highlight: matched nodes plus their immediate neighbors (parents and children).
    highlight_nodes = set()
    for node in search_results:
        highlight_nodes.add(node)
        highlight_nodes.update(G.predecessors(node))
        highlight_nodes.update(G.successors(node))

    # Create the graph layout and plot.
    pos = nx.spring_layout(G, k=0.5, iterations=50)
    plt.figure(figsize=(12, 8))
    node_colors = ['pink' if node in highlight_nodes else 'lightblue' for node in G.nodes]
    nx.draw(G, pos, with_labels=True, node_size=1500, node_color=node_colors, font_size=10, arrows=True)
    plt.title("Ontology Graph with Highlighted Related Nodes\n" + result_message)

    # Save the plot to a temporary file and close the plot.
    tmpfile = tempfile.NamedTemporaryFile(suffix='.png', delete=False)
    plt.savefig(tmpfile.name, format="png")
    plt.close()

    return tmpfile.name, result_message

# Load default file and search term
default_owl_path = "genai.owl"
default_search_term = "LoRa"
with open(default_owl_path, "rb") as f:
    default_owl_bytes = f.read()

# Create a temporary file to store the default OWL content
temp_owl_file = tempfile.NamedTemporaryFile(delete=False, suffix=".owl")
temp_owl_file.write(default_owl_bytes)
temp_owl_file.close()

# Define the Gradio interface.
iface = gr.Interface(
    fn=plot_ontology,
    inputs=[
        gr.Textbox(lines=1, label="Search Term", placeholder="Enter a search term (e.g., LoRa)", value=default_search_term),
        gr.File(label="Upload Ontology OWL File", value=temp_owl_file.name)
    ],
    outputs=[
        gr.Image(type="filepath", label="Ontology Graph"),
        gr.Textbox(label="Search Result Details")
    ],
    title="Ontology Explorer",
    description="Upload an ontology OWL file and enter a search term to visualize the ontology graph. "
                "Nodes matching the search term and their immediate neighbors (parents and children) are highlighted in pink."
                "Reference of the Ontology Loaded: [GENAI](https://bioportal.bioontology.org/ontologies/GENAI)",
    examples=[
        ["LoRa", "genai.owl"],
        ["LLM", "genai.owl"],
        ["RLHF", "genai.owl"]
    ],
    cache_examples=True
)

if __name__ == "__main__":
    iface.launch(share=True)