import gradio as gr
import pandas as pd
import spacy
import concepcy

LIST_RELATIONS = ["RelatedTo", "FormOf", "IsA", "PartOf", "HasA", "UserFor", "CapableOf", "AtLocation", "Causes",
                  "HasSubevent", "HasFirstSubevent", "HasLastSubevent", "HasPrerequisite", "HasProperty",
                  "MotivatedByGoal", "ObstructedBy", "Desires", "CreatedBy", "Synonym", "Antonym", "DistinctFrom",
                  "DerivedFrom", "SymbolOf", "DefinedAs", "MannerOf", "LocatedNear", "HasContext", "SimilarTo",
                  "EtymologicallyRelatedTo", "EtymologicallyDerivedFrom", "CausesDesire", "MadeOf", "ReceivesAction",
                  "ExternalURL"]

nlp = spacy.load("en_core_web_sm")
nlp.add_pipe(
    "concepcy",
    config={
        "relations_of_interest": LIST_RELATIONS,
        "filter_missing_text": True,
        "filter_edge_weight": 2,
    }
)


def greet(text, relations, to_enrich):
    doc = nlp(text)

    outs0, outs1 = [], []
    if "document" in to_enrich:
        outs0 = []
        for relation in relations:
            for r in doc._.get(relation.lower()).values():
                outs0.extend([[relation, elt["text"]] for elt in r])

    if "token" in to_enrich:

        for token in doc:
            for relation in relations:
                rels = token._.get(relation.lower())
                if len(rels) > 0:
                    print(rels)
                    print([[token.text, relation, r["text"]] for r in rels])
                    outs1.extend([[token.text, relation, r["text"]] for r in rels])

    return pd.DataFrame(outs0, columns=["relation", "text"]), pd.DataFrame(outs1, columns=["word", "relation", "text"])


iface = gr.Interface(
    fn=greet,
    title="Playground for <a href='https://github.com/JulesBelveze/concepcy'>concepCy</a>",
    description="This demo enables you to play around with SpaCy's concepCy wrapper, a wrapper for ConceptNet!"
                "To get started: enter a piece of text, check the relations you are interested in and if you want to "
                "retrieve relations at a document-level and/or token-level.\n"
                "Relations will be displayed as tables on the right hand side! Have fun!\n",
    inputs=[
        gr.Textbox(placeholder="Enter sentence here...", lines=5, value="I love eating pizzas"),
        gr.CheckboxGroup(choices=LIST_RELATIONS, value=["IsA"]),
        gr.CheckboxGroup(choices=["document", "token"], value=["document"])
    ],
    outputs=[
        gr.Dataframe(headers=["relation", "text"], label="Document-level relations"),
        gr.Dataframe(headers=["word", "relation", "text"], label="Token-level relations")
    ]
)
iface.launch()