import streamlit as st
import openai
import pyvis.network as net
from streamlit.components.v1 import html

# Set your OpenAI API key
openai.api_key = "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"

# Function to generate ideas using GPT
def generate_ideas(prompt):
    response = openai(
        engine="text-davinci-003",
        prompt=f"Generate related ideas for: {prompt}",
        max_tokens=100
    )
    return response.choices[0].text.strip().split("\n")

# Function to create a mind map using pyvis
def create_mind_map(main_idea, ideas):
    g = net.Network(height="500px", width="100%", directed=True)
    g.add_node(0, label=main_idea, color="red", size=30)
    
    for i, idea in enumerate(ideas):
        g.add_node(i+1, label=idea, color="blue", size=20)
        g.add_edge(0, i+1)

    return g

# Streamlit app UI
st.title("Mind Mapping with AI Idea Expander")

# Input for main idea
main_idea = st.text_input("Enter your main idea:")

if st.button("Generate Mind Map"):
    if main_idea:
        # Generate related ideas
        related_ideas = generate_ideas(main_idea)
        st.write("Related Ideas:")
        for idea in related_ideas:
            st.write(f"- {idea}")
        
        # Create mind map
        g = create_mind_map(main_idea, related_ideas)
        
        # Save mind map as HTML and display it
        g.save_graph("mind_map.html")
        with open("mind_map.html", "r", encoding="utf-8") as f:
            html(f.read(), height=600)
    else:
        st.error("Please enter an idea.")