import streamlit as st
import os
import networkx as nx
from streamlit_agraph import agraph, Node, Edge, Config
import pandas as pd
import numpy as np
from sklearn.metrics.pairwise import cosine_similarity, manhattan_distances, euclidean_distances


# Heading
st.title('Clinical Trial Knowledge Graph Demo [Eye Diseases]')

mapping = np.load("mapping.npy", allow_pickle=True).item()
conditions = pd.read_csv("condition_count.csv")['condition'].values.tolist()
all_graphs = np.load("by_condition_graph.npy", allow_pickle=True).item()
emb_data = np.load("emb_data.npy", allow_pickle=True).item()

def get_all():
    nodes, edges = [], []
    for c in conditions:
        G = all_graphs[c]
        nodes.append([Node(id=i, title=i, size=20) for i in G.nodes])
        edges.append([Edge(source=i, target=j, weight=weight,   type="CURVE_SMOOTH") for (i,j, weight) in G.edges.data()])

    nodes = [item for sublist in nodes for item in sublist]
    edges = [item for sublist in edges for item in sublist]


    return nodes, edges

        
        


option = st.selectbox(
     'Select a condition',
   conditions + ["All"])

if option == 'All':
    nodes, edges = get_all()
else:
    G = all_graphs[option]
    nodes = [Node(id=i, title=i, size=20) for i in G.nodes]
    st.info(f"Numer of nodes: {len(nodes)}")
    edges = [Edge(source=i, title="Similarity is " +  str(weight['weight']), target=j, weight=weight['weight'],   type="CURVE_SMOOTH") for (i,j, weight) in G.edges.data()]

config = Config(width=700, 
                height=500, 
                directed=False,
                nodeHighlightBehavior=True, 
                highlightColor="#F7A7A6",
                collapsible=True,
                initialZoom=1,
                node={'labelProperty':'label'},
                link={'labelProperty': 'label', 'renderLabel': True}
                ) 

return_value = agraph(nodes=nodes, 
                      edges=edges, 
                      config=config)