from pyvis.network import Network
import streamlit as st

def explain(answer):
    graph = Network(height="450px", width="100%")
    sources_table = []
            #all_reference_texts = ''
    for nodewithscore in answer.source_nodes:  
        node = nodewithscore.node
        from llama_index.core.schema import NodeRelationship
        if NodeRelationship.SOURCE in node.relationships:
            node_id = node.relationships[NodeRelationship.SOURCE].node_id
            node_id = node_id.split('/')[-1]
            title = node_id.split('.')[2].replace('_', ' ')
            link = '.'.join(node_id.split('.')[:2])[:10]
            link = f'https://arxiv.org/abs/{link}'
            href = f'<a  target="_blank" href="{link}">{title}</a>'
            sources_table.extend([[href, node.text]])
                    #all_reference_texts = all_reference_texts + '\n' + node.text
        else:
                    #st.write(node.text) TODO second level relationships
            rel_map = node.metadata['kg_rel_map']
            for concept in rel_map.keys():
                        #st.write(concept)
                graph.add_node(concept,  concept, title=concept)
                rels = rel_map[concept]
                for rel in rels:
                    graph.add_node(rel[1],  rel[1], title=rel[1])
                    graph.add_edge(concept,  rel[1], title=rel[0])
            # --- display the query terms graph
    st.session_state.graph_name = 'graph.html'
    graph.save_graph(st.session_state.graph_name)            
    import streamlit.components.v1 as components
    graphHtml = open(st.session_state.graph_name, 'r', encoding='utf-8')
    source_code = graphHtml.read() 
    components.html(source_code, height = 500)
            # --- display the reference texts table
    import pandas as pd
    df = pd.DataFrame(sources_table)
    df.columns = ['paper', 'relevant text']
    st.markdown(""" <style> 
                                table[class*="dataframe"] {
                                font-size: 10px;
                            }
                            </style> """, unsafe_allow_html=True)
    st.write(df.to_html(escape=False), unsafe_allow_html=True)