import gradio as gr
from anytree import Node, RenderTree
import spacy
from spacy import displacy

nlp = spacy.load("en_core_web_md")

# -- Main functions --


def postag(tk):
    tag = ""
    plural_tags = ["NNS", "NNPS"]
    if tk.tag_ in plural_tags:
        tag = " (Plural)"
    else:
        tag = " ({})".format(tk.tag_)
    return tag
    
def genFlatDepthTree(expr):
    doc = nlp(expr)
    root = next(doc.sents).root
    node = Node(root.text+"\t (Root)"+postag(root)+" - [0]", parent=None)

    def tree(tk, last_node, depth):
        if tk.n_lefts + tk.n_rights > 0:
            for ch in tk.children:
                tree(ch, Node(ch.text+"\t"+postag(ch)+" - ["+str(depth+1)+"]", parent=last_node), depth+1)
    
    tree(root, node, 0)
    flat_tree = ""

    for pre, fill, node in RenderTree(node):
        flat_tree += """{}{}\n""".format(pre, node.name)

    img_tree = displacy.render(doc, style='dep', options={'distance': 100})
    
    return "<center><div style='max-width: 800px; overflow-x:auto;'>"+img_tree+"</div></center>", flat_tree
    
# def syntacticTree(expr):
#     doc = nlp(expr)
#     depths = {}

#     def walk_tree(tk, depth):
#         depths[tk.text] = depth
#         if tk.n_lefts + tk.n_rights > 0: # if Doesn't root
#             [walk_tree(child, depth + 1) for child in tk.children]

#     [walk_tree(sent.root, 0) for sent in doc.sents]

#     flat_tree = [tk.text+" - "+str(depths[tk.text]) for tk in doc]
#     image_tree = displacy.render(doc, style='dep', options={'distance': 100})
#     image_tree = "<center><div style='max-width: 800px; overflow-x:auto;'>"+image_tree+"</div></center>"
    
#     return image_tree,flat_tree

# -- Interface --
demo = gr.Blocks(css=".container { max-width: 900px; margin: auto;}", title="Syntactic tree generator")

with demo:
    gr.Markdown("""<center> <h2>🌳 Syntactic Tree Generator 🌳</h2> </center>""")
    with gr.Row():
        with gr.Column():
            gr.Image(value="https://img.unocero.com/2019/11/facebook-app-para-hacer-memes-1-1024x576.jpg",label="", type="URL")
        with gr.Column():
            input_expr = gr.Textbox(label="Input", placeholder="Enter an expression here")
            out_flat_tree = gr.Text(label="Flat tree", value="")
    
    out_image_tree = gr.HTML(label="")

    gr.Examples(
        examples=[ 
            "glass in guys hand on right first guy",
            "i have been happy",
            "the happy spiders",
            "the best dog out there",
             "girl 's face",
             'bottom grass',
             'rock people are sitting on',
             'blue sky center below clouds',
             'group of people on left',
             'tree middle',
             'the lump of grass above the bright rock in the bottom left',
             'berries in middle',
             'red shirt kid',
             'middle rock',
             'grass below the person in the straw hat',
             'grass on the left',
             'wall between stairs 2nd lv on right',
             'the large group of clouds',
             'sky top left',
             'rock structure on far right',
             'left donkey'],
    inputs=input_expr,
        examples_per_page=5,
    )
    
    input_expr.change(     
        fn=genFlatDepthTree, 
        inputs=input_expr, 
        outputs=[out_image_tree, out_flat_tree])

demo.launch()