import os
import networkx as nx
import matplotlib.pyplot as plt
from langchain_groq import ChatGroq
from langchain.chains import LLMChain
from langchain.prompts import PromptTemplate
import gradio as gr

# Set up the ChatGroq API
GROQ_API_KEY = os.environ.get("GROQ_API_KEY")
if not GROQ_API_KEY:
    raise ValueError("Please set the GROQ_API_KEY environment variable")

# Initialize the ChatGroq LLM
llm = ChatGroq(temperature=0, model_name='llama-3.1-8b-instant', groq_api_key=GROQ_API_KEY)

# Define a prompt template for entity extraction
entity_extraction_prompt = PromptTemplate(
    input_variables=["text"],
    template="Extract the main entities and their relationships from the following text:\n{text}\n\nEntities and relationships:"
)

# Create an LLMChain for entity extraction
entity_chain = LLMChain(llm=llm, prompt=entity_extraction_prompt)

def create_knowledge_graph(text):
    """Extracts entities and relationships from text and creates a knowledge graph."""
    result = entity_chain.run(text)

    # Parse the result and create a graph
    G = nx.Graph()

    # Simple parsing logic (you may need to adjust this based on the LLM's output format)
    lines = result.split('\n')
    for line in lines:
        if '-' in line:
            entity1, rest = line.split('-', 1)
            relation, entity2 = rest.split(':', 1)
            entity1 = entity1.strip()
            entity2 = entity2.strip()
            relation = relation.strip()

            G.add_node(entity1)
            G.add_node(entity2)
            G.add_edge(entity1, entity2, relationship=relation)

    return G

def visualize_graph(G):
    """Visualizes the knowledge graph using matplotlib."""
    pos = nx.spring_layout(G)
    plt.figure(figsize=(12, 8))
    nx.draw(G, pos, with_labels=True, node_color='lightblue', node_size=3000, font_size=10, font_weight='bold')
    edge_labels = nx.get_edge_attributes(G, 'relationship')
    nx.draw_networkx_edge_labels(G, pos, edge_labels=edge_labels)
    plt.title("Knowledge Graph")
    plt.axis('off')
    plt.tight_layout()
    return plt

def generate_knowledge_graph(text):
    """Generates and visualizes a knowledge graph from input text."""
    if text:
        knowledge_graph = create_knowledge_graph(text)
        fig = visualize_graph(knowledge_graph)
        return fig
    else:
        return None

# Custom footer HTML
footer_html = """
<footer>
    <p>If you enjoyed the functionality of the app, please leave a like!<br>
    Check out more on <a href="https://www.linkedin.com/in/girish-wangikar/" target="_blank">LinkedIn</a> |
    <a href="https://girishwangikar.github.io/Girish_Wangikar_Portfolio.github.io/" target="_blank">Portfolio</a></p>
</footer>
"""

# Instructions for using the app
instructions_html = """
<div>
    <h2>Instructions:</h2>
    <ol>
        <li>Enter text in the input box that contains entities and their relationships. For example, "Paris - capital of: France".</li>
        <li>Click the "Submit" button to generate the knowledge graph.</li>
        <li>View the resulting knowledge graph visualization, which will display the entities as nodes and their relationships as labeled edges.</li>
        <li>Feel free to experiment with different texts to explore relationships visually!</li>
    </ol>
</div>
"""

# Gradio interface
iface = gr.Interface(
    fn=generate_knowledge_graph,
    inputs=gr.Textbox(lines=10, placeholder="Enter your text here..."),
    outputs=gr.Plot(),
    title="Knowledge Graph Generator",
    description="Enter text to generate a knowledge graph.",
    article=instructions_html + footer_html,
    theme="default",
    css="""
    footer {
        margin-top: 20px;
        text-align: center;
        color: #bb86fc;
        position: fixed;
        bottom: 0;
        width: 100%;
        background-color: white;
        padding: 10px 0;
    }
    footer a {
        color: #bb86fc !important;
        text-decoration: none;
    }
    footer a:hover {
        text-decoration: underline;
    }
    """)

# Launch the interface
if __name__ == "__main__":
    iface.launch()