import streamlit as st
from streamlit_agraph import agraph, Node, Edge, Config

from meutils.pipe import *

Node = lru_cache(100000)(Node)
Edge = lru_cache(100000)(Edge)


def df2nodes(df_nodes):  # df.itertuples()
    df_nodes.columns = ['k', 'v', 'r']

    nodes = set()
    edges = set()

    for source, target, label in tqdm(df_nodes.values, 'Triplet'):
        node_s = Node(source, label=source)
        node_t = Node(target, label=target)

        edge = Edge(source, target, label=label)

        nodes.update([node_s, node_t])
        edges.update([edge])
    return nodes, edges


# df = pd.read_excel("~/Desktop/simple.xlsx", sheet_name="links")
df = pd.read_csv('~/Desktop/tt.csv').tail(500)


nodes, edges = df2nodes(df)

config = Config(width=1000, height=1000)


st.set_page_config(page_title="Knowledge Graph", page_icon="💗", layout="wide")
return_value = agraph(nodes=nodes, edges=edges, config=config)

del df, nodes, edges