from huggingface_hub import InferenceClient
import ast
import nltk
import matplotlib.pyplot as plt
import gradio as gr

client = InferenceClient("Qwen/Qwen2.5-72B-Instruct")


def get_structures(sent):
    c_structure = ast.literal_eval(client.chat.completions.create(
        messages=[
            {"role": "system",
             "content": "generate bnf description for buiding c-structure according to lexical-functional grammar framework, no explanation or additional text, use the following structure:\n"
                        "c-structure: 'generated bnf description'"
             },
            {"role": "user",
             "content": f"generate bnf description for c-structure of the following sentence: {sent}"},
        ],
        response_format={
            "type": "json",
            "value": {
                "properties": {
                    "c-structure": {"type": "string"}},
            }
        },
        stream=False,
        max_tokens=512,
        temperature=0.7,
        top_p=0.1
    ).choices[0].get('message')['content'])

    c_latex = ast.literal_eval(client.chat.completions.create(
        messages=[
            {"role": "system",
             "content": "generate nltk respresentation for the LFG c-structure of the sentence according to provided bnf description, no explanation or additional text\n"
                        "example: (S (NP 'Text') (VP 'text')))"
             },
            {"role": "user",
             "content": f"description: {c_structure['c-structure']}"},
        ],
        response_format={
            "type": "json",
            "value": {
                "properties": {
                    "c-structure": {"type": "string"}},
            }
        },
        stream=False,
        max_tokens=512,
        temperature=0.7,
        top_p=0.1
    ).choices[0].get('message')['content'])

    tree = nltk.Tree.fromstring(c_latex['c-structure'])

    with open('output.txt', 'wt') as out:
        tree.pretty_print(stream=out)

    with open('output.txt', 'a') as f:
        f.write(f'c-structure:\n{c_latex["c-structure"]}\n\nBNF-description:\n{c_structure["c-structure"]}')

    return 'output.txt'

interface = gr.Interface(
    fn=get_structures,
    inputs=gr.Textbox(label="Enter your sentence"),
    outputs=gr.File(),
    title="LFG AI-Parser",
    description="Enter a sentence and visualize its c-structure according to LFG.",
)

if __name__ == "__main__":
    interface.launch(share=True)