import gradio as gr
import pandas as pd
from nltk.corpus import wordnet as wn
import nltk
nltk.download('wordnet')

df_tree = pd.read_csv("data/20k_tree.csv")
df_metrics = pd.read_csv("data/benchmark.csv")

def query_interface(model, synset):
    if model == "*":
        return df_metrics.query(f"synset =='{synset}'")
    else:
        return df_metrics.query(f"synset =='{synset}' and model == '{model}'")    

def find_synsets(lemma):
    synsets = wn.synsets(lemma)
    valid_set = set(df_metrics.synset.values)
    synset_names = [s.name() for s in synsets if s.name() in valid_set]
    return df_tree.query(f"synset in @synset_names", engine='python')
    
with gr.Blocks() as app:
    gr.Label("WordNet Explorer", style="h1")
    lemma = gr.Textbox(label="find synsets")
    gr.Button("Search").click(find_synsets, lemma, outputs=[gr.DataFrame()])
    
    gr.Label("Benchmark Explorer", style="h1")
    model_list = list(set(df_metrics.model.values.tolist()))
    gr.TextArea(f"This is a simple interface to query the benchmark results. You can use * as a wildcard for the model name. The available models are {', '.join(model_list)}.", lines=2)
    with gr.Row():
        # model = gr.inputs.Dropdown(label="model",choices=list(set(df_metrics.model.values.tolist())))
        model = gr.Textbox(label="model")
        synset = gr.inputs.Textbox(label="synset")
    btn = gr.Button("Query", label="query")
    btn.click(query_interface, inputs=[model, synset], outputs=[gr.DataFrame()])
    
app.launch()