import gradio as gr
import networkx as nx
from pyvis.network import Network
import spacy
import spacy.cli
import torch
import json
from transformers import pipeline
import tempfile
import os

class KnowledgeGraphVisualizer:
    def __init__(self):
        # Download and load spaCy model if not already available
        try:
            self.nlp = spacy.load("en_core_web_sm")
        except OSError:
            spacy.cli.download("en_core_web_sm")  # Download the model
            self.nlp = spacy.load("en_core_web_sm")  # Load the model after downloading

        # Initialize summarization pipeline
        self.summarizer = pipeline("summarization", device=0 if torch.cuda.is_available() else -1)
        
        # Initialize the graph
        self.G = nx.DiGraph()
    
    def extract_triples(self, text):
        """Extract subject-predicate-object triples from text using spaCy"""
        doc = self.nlp(text)
        triples = []
        
        for sent in doc.sents:
            for token in sent:
                if token.dep_ == "ROOT":
                    subject = ' '.join([tok.text for tok in token.lefts if tok.dep_ in ["nsubj", "nsubjpass"]])
                    obj = ' '.join([tok.text for tok in token.rights if tok.dep_ in ["dobj", "pobj"]])
                    if subject and obj:
                        triples.append((subject, token.text, obj))
        
        print(f"Extracted triples: {triples}")  # Debugging
        return triples
    
    def add_triples_to_graph(self, triples):
        """Add triples to the graph"""
        for subj, pred, obj in triples:
            self.G.add_edge(subj, obj, label=pred)
    
    def visualize_graph(self):
        """Create an HTML visualization of the graph"""
        net = Network(notebook=False, directed=True, height="600px", width="100%")
        
        # Add nodes and edges from the NetworkX graph
        for node in self.G.nodes():
            net.add_node(node, label=node)
        
        for edge in self.G.edges(data=True):
            net.add_edge(edge[0], edge[1], label=edge[2]['label'])
        
        # Generate temporary file
        temp_dir = tempfile.mkdtemp()
        path = os.path.join(temp_dir, "graph.html")
        net.save_graph(path)
        
        return path
    
    def summarize_graph(self):
        """Generate a summary of the graph structure"""
        nodes = list(self.G.nodes())
        edges = list(self.G.edges(data=True))
        
        summary_text = f"The knowledge graph contains {len(nodes)} nodes and {len(edges)} relationships. "
        
        # Create a natural language description of some relationships
        if edges:
            summary_text += "Some key relationships include: "
            for i, (src, dst, data) in enumerate(edges[:3]):
                summary_text += f"{src} {data['label']} {dst}. "
        
        summary = self.summarizer(summary_text, max_length=100, min_length=30)[0]['summary_text']
        print(f"Generated summary: {summary}")  # Debugging
        return summary
    
    def clear_graph(self):
        """Clear the current graph"""
        self.G.clear()
        return self.visualize_graph()

def create_gradio_interface():
    # Initialize the visualizer
    visualizer = KnowledgeGraphVisualizer()
    
    def process_input(text, input_type):
        if input_type == "Text":
            triples = visualizer.extract_triples(text)
        else:
            try:
                # Assume JSON format: [["subject", "predicate", "object"], ...]
                triples = json.loads(text)
            except:
                return "Error: Invalid JSON format", "Error: Could not process input"
        
        visualizer.add_triples_to_graph(triples)
        summary = visualizer.summarize_graph()
        graph_path = visualizer.visualize_graph()
        
        # Return the HTML file path and summary
        return graph_path, summary
    
    def clear():
        graph_path = visualizer.clear_graph()
        return graph_path, ""
    
    # Create the interface
    with gr.Blocks() as demo:
        gr.Markdown("# Knowledge Graph Visualizer")
        
        with gr.Row():
            with gr.Column():
                input_text = gr.Textbox(label="Input Data", lines=5)
                input_type = gr.Radio(["Text", "JSON"], label="Input Type", value="Text")
                
                with gr.Row():
                    submit_btn = gr.Button("Process")
                    clear_btn = gr.Button("Clear Graph")
            
            with gr.Column():
                output_viz = gr.File(label="Graph Visualization", file_types=[".html"])
                output_summary = gr.Textbox(label="Graph Summary", lines=3)
        
        submit_btn.click(
            fn=process_input,
            inputs=[input_text, input_type],
            outputs=[output_viz, output_summary]
        )
        
        clear_btn.click(
            fn=clear,
            inputs=[],
            outputs=[output_viz, output_summary]
        )
    
    return demo

# Launch the interface
if __name__ == "__main__":
    demo = create_gradio_interface()
    demo.launch(share=True)