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

title = "RE Generation"
description_GNER = """
A T5 model tuned to perform text generation, and predict which characters are present in the report. Note that, in the Hall and Van de Castle, the character lists never includes the dreamer. Hence, if you (willingly or not) enter a report that does not contain another character reference, the model will/should (correctly) produce an empty string. Moreover, it is likely that the produced list of CHAR could be longer than the one produced by the SA model, as not all CHAR might be associated with emotions. 
"""

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, 
        token=usr_tkn,
        max_length=300,
        min_length=10,
    )

    # 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,
# )