import streamlit as st
import streamlit.components.v1 as components
from pyvis.network import Network
import networkx as nx
import re
import emoji

st.title('Pyvis VisJS Knowledge Graph from Markdown')

# Make Network show itself with repr_html
def net_repr_html(self):
    nodes, edges, height, width, options = self.get_network_data()
    html = self.template.render(height=height, width=width, nodes=nodes, edges=edges, options=options)
    return html

Network._repr_html_ = net_repr_html

st.sidebar.title('Upload a Markdown File')
uploaded_file = st.sidebar.file_uploader("Choose a file", type=["md"])

physics = st.sidebar.checkbox('Add physics interactivity?')

def draw_knowledge_graph(file_contents, physics):
    # Extract sentences and emojis from the file
    sentences = re.split(r'\n', file_contents)
    emojis_list = [emoji.emojize(x) for x in emoji.UNICODE_EMOJI['en']]
    emojis = [ch for ch in file_contents if ch in emojis_list]
    
    # Create a graph
    G = nx.Graph()

    # Add nodes and edges
    for idx, sentence in enumerate(sentences):
        G.add_node(idx, label=sentence)
        words = set(sentence.lower().split())
        for j, other_sentence in enumerate(sentences[:idx]):
            other_words = set(other_sentence.lower().split())
            common_words = words & other_words
            common_emojis = set(emojis) & set(sentence) & set(other_sentence)
            if common_words:  # or common_emojis
                G.add_edge(idx, j, weight=len(common_words) + len(common_emojis))
    
    # Visualize the graph
    net = Network(notebook=True)
    net.from_nx(G)
    net.toggle_physics(physics)

    return net._repr_html_()

if uploaded_file is not None:
    file_contents = uploaded_file.read().decode("utf-8")
    html = draw_knowledge_graph(file_contents, physics)
    components.html(html, height=1200, width=1000)
else:
    st.warning("Please upload a markdown file to visualize the knowledge graph.")