from git import Repo
import gradio as gr
from huggingface_hub import snapshot_download
import penman
import sys

Repo.clone_from("https://github.com/AbdiHaryadi/amr-tst-indo.git", "amr-tst-indo")
sys.path.append("./amr-tst-indo")

from text_to_amr import TextToAMR

# Gold v2
amr_parsing_model_name = "mbart-en-id-smaller-indo-amr-parsing-translated-nafkhan"
snapshot_download(
    repo_id=f"abdiharyadi/{amr_parsing_model_name}",
    local_dir=f"./amr-tst-indo/AMRBART-id/models/{amr_parsing_model_name}",
    ignore_patterns=[
        "*log*",
        "*checkpoint*",
    ]
)
t2a = TextToAMR(model_name=amr_parsing_model_name)

def run(text, source_style):
    source_amr, *_ = t2a([text])
    source_amr_display = penman.encode(source_amr)
    yield source_amr_display, "...", "...", "...", "..."

    triplets = [
        ("kamar", "sangat bagus", "positif"),
        ("kamar", "bersih", "positif")
    ]
    triplets_display = "\n".join(f"({x[0]}, {x[1]}, {x[2]})" for x in triplets)
    yield source_amr_display, triplets_display, "...", "...", "..."

    style_words = ["bagus", "bersih"]
    style_words_display = ", ".join(style_words)
    yield source_amr_display, triplets_display, style_words_display, "...", "..."

    target_amr = penman.decode("(z0 / dunia)")
    target_amr_display = penman.encode(target_amr)
    yield source_amr_display, triplets_display, style_words_display, target_amr_display, "..."

    result = f"dunia ({text=}, {source_style=})"
    yield source_amr_display, triplets_display, style_words_display, target_amr_display, result

demo = gr.Interface(
    fn=run,
    inputs=[
        gr.Textbox(label="Teks (Text)"),
        gr.Radio(label="Gaya sumber (Source style)", choices=[
            ("Positif (Positive)", "LABEL_1"),
            ("Negatif (Negative)", "LABEL_0"),
        ], value="LABEL_1"),
    ],
    outputs=[
        gr.Textbox(label="Graf AMR sumber (Source AMR graph)"),
        gr.Textbox(label="Triplet (Triplets)"),
        gr.Textbox(label="Kata bergaya (Style words)"),
        gr.Textbox(label="Graf AMR target (Target AMR graph)"),
        gr.Textbox(label="Hasil (Result)"),
    ]
)
demo.launch()