from pprint import pprint, pformat
import gradio as gr
import click
from rasa.nlu.model import Interpreter


MODEL_PATH = "woz_nlu_agent/models/nlu"
interpreter = None

def predict(input):

    def rasa_output(text):
        message = str(text).strip()
        result = interpreter.parse(message)
        return result

    response = rasa_output(input)
    
    del response["response_selector"]
    response["intent_ranking"] = response["intent_ranking"][:3]
    if "id" in ressponse["intent"]:
        del response["intent"]["id"]
    for i in response["intent_ranking"]:
        if "id" in i:
            del i["id"]
    for e in response["entities"]:
        if "extractor" in e:
            del e["extractor"]
        if "start" in e and "end" in e:
            del e["start"]
            del e["end"]

    return pformat(response, indent=4)


def main():
    global interpreter
    print("Loading model...")
    import os
    print(os.listdir("woz_nlu_agent/models/nlu"))
    interpreter = Interpreter.load(MODEL_PATH)
    print("Model loaded.")
    iface = gr.Interface(fn=predict, inputs="text", outputs="text")
    iface.launch()


if __name__ == "__main__":
    main()