import spacy
import gradio as gr
from spacy import displacy

# Load your trained spaCy model
nlp = spacy.load("sr_Spacy_Serbian_Model_SrpKor4Tagging_BERTICOVO")

# Define a function to display the tags and lemmas
def display_tags_and_lemmas(text):
    doc = nlp(text)
    html = displacy.render(doc, style="ent", page=True, minify=True)

    # We'll also create a custom HTML to display lemmas nicely
    lemma_html = "<div class='lemma-display'><table class='table table-hover'>"
    lemma_html += "<thead><tr><th>Token</th><th>Lemma</th><th>POS Tag</th></tr></thead><tbody>"

    for token in doc:
        lemma_html += f"<tr><td>{token.text}</td><td>{token.lemma_}</td><td>{token.pos_}</td></tr>"

    lemma_html += "</tbody></table></div>"

    # Return both the displaCy HTML and our custom lemma table
    return html, lemma_html

# Define Gradio interface
iface = gr.Interface(
    fn=display_tags_and_lemmas,
    inputs=gr.inputs.Textbox(lines=5, placeholder="Unesite rečenicu ovde..."),
    outputs=[
        gr.outputs.HTML(label="Leme i POS oznake")
    ],
    title="spaCy Tagger i Lemmatizer",
    description="Unesite rečenicu da biste videli njene imenovane entitete, POS oznake i leme.",
    examples=["Lep dan, danas."]
)


# Include CSS styling for nicer outputs
css = """
    <style>
        .lemma-display { margin: auto; width: 50%; }
        .table { font-family: Arial, Helvetica, sans-serif; border-collapse: collapse; width: 100%; }
        .table td, .table th { border: 1px solid #ddd; padding: 8px; }
        .table tr:nth-child(even){ background-color: #f2f2f2; }
        .table tr:hover { background-color: #ddd; }
        .table th { padding-top: 12px; padding-bottom: 12px; text-align: left; background-color: #04AA6D; color: white; }
    </style>
"""

iface.launch(inline=False, enable_queue=True, share=True, css=css)