from pyvis.network import Network
import gradio as gr
from transformers import pipeline
import os

model_id = "models/DReAMy-lib/t5-base-DreamBank-Generation-Act-Char"

def get_graph_dict(graph_text):
    edge_labels = {}
    if graph_text == "":
        edge_labels = {("No_Graphs", None):None}

    else:
        try:
            for trpl in graph_text[1:-1].split(" | "):
                h,r,t = trpl[1:-1].split(" # ") 
                edge_labels[(h,t)] = r
        except:
            edge_labels = {("Error", None):None}
    return edge_labels

def text_to_graph(text):
    
    # Use a pipeline as a high-level helper
    pipe = pipeline(
        "text2text-generation", 
        model=model_id, 
        max_length=300,
        min_length=5,
    )

    # generate text graph
    graph_text = pipe(text)

    graph_text = graph_text[0]["generated_text"]

    # get the nodes: label dict
    edge_labels = get_graph_dict(graph_text)
    
    # create network
    net = Network(directed=True)

    # nodes & edges
    for (h, t), r  in edge_labels.items():
        if (h == "Error") or (h == "No_Graphs"):
            net.add_node(h, shape="circle")
            continue
        else:
            net.add_node(h, shape="circle")
            net.add_node(t, shape="circle")
            net.add_edge(h, t, title=r, label=r)

    # set structure
    net.repulsion(
        node_distance=200,
        central_gravity=0.2,
        spring_length=200,
        spring_strength=0.05,
        damping=0.09
    )

    net.set_edge_smooth('dynamic') 

    # get html
    html = net.generate_html()
    html = html.replace("'", "\"")
    
    html_s = f"""<iframe style="width: 100%; height: 600px;margin:0 auto" name="result" allow="midi; geolocation; microphone; camera; 
    display-capture; encrypted-media;" sandbox="allow-modals allow-forms 
    allow-scripts allow-same-origin allow-popups 
    allow-top-navigation-by-user-activation allow-downloads" allowfullscreen="" 
    allowpaymentrequest="" frameborder="0" srcdoc='{html}'></iframe>"""

    return html_s, graph_text

# demo = gr.Interface(
#         text_to_graph,
#         inputs=gr.Textbox(label="Text", placeholder="Enter a text here."),
#         outputs=[gr.HTML(label="Extracted graph"),gr.Textbox(label="Extracted text")]
#         examples= [
#         ["I was skating on the outdoor ice pond that used to be across the street from my house. I was not alone, but I did not recognize any of the other people who were skating around. I went through my whole repertoire of jumps, spires, and steps-some of which I can do and some of which I'm not yet sure of. They were all executed flawlessly-some I repeated, some I did only once. I seemed to know that if I went into competition, I would be sure of coming in third because there were only three contestants. Up to that time I hadn't considered it because I hadn't thought I was good enough, but now since everything was going so well, I decided to enter."],
#         ["I was talking on the telephone to the father of an old friend of mine (boy, 21 years old). We were discussing the party the Saturday night before to which I had invited his son as a guest. I asked him if his son had a good time at the party. He told me not to tell his son that he had told me, but that he had had a good time, except he was a little surprised that I had acted the way I did."], 
#         ["I was walking alone with my dog in a forest."]
#         ],
#         title=title,
#         description=description_GNER,
#         cache_examples=True,
# )