import streamlit as st
from flair.data import Sentence
from flair.models import SequenceTagger

# Load the Flair model
model_path = "onurkeles/hamshetsnag-pos-tagger"
pos_tagger = SequenceTagger.load(model_path)

def tag_pos(text, detailed_output):
    """Tag parts of speech in a given text, with optional detailed output."""
    sentence = Sentence(text)
    pos_tagger.predict(sentence)

    if detailed_output:
        # Generate detailed information with tag values and probabilities
        return "\n".join(
            [f"{token.text}: {token.get_tag('pos').value} ({token.get_tag('pos').score:.2f})" for token in sentence]
        )
    else:
        # Return a simple tagged string
        return sentence.to_tagged_string()

def write():
    st.markdown("# Part-of-Speech Tagging for Hamshetsnag")
    st.sidebar.header("POS Tagging")

    st.write(
        '''Detect parts of speech in Hamshetsnag text using the fine-tuned model.'''
    )

    # Sidebar for configurations
    st.sidebar.subheader("Configurable Parameters")

    # Detailed Output Checkbox
    detailed_output = st.sidebar.checkbox(
        "Detailed Output",
        value=False,
        help="If checked, output shows detailed tag information (probability scores, etc.).",
    )

    # Input field for text
    input_text = st.text_area(label='Enter a text: ', height=100, value="Put example text here.")
    
    # Provide example sentences with translations
    example_sentences = [
        "tuute acertsetser topoldetser. aaav ta? (TR: Kâğıdı büzüştürdün attın. Oldu mu?)",
        "Baran u Baden teran. (TR: Baran ve Bade koştu.)",
        "Onurun ennush nu İremin terchushe intzi shad kızdırmısh aaav. (TR: Onur'un düşüşü ve İrem'in koşuşu beni kızdırdı.)"
    ]

    st.write("## Example Sentences:")
    for example in example_sentences:
        if st.button(f"Use: {example.split('(TR:')[0].strip()}"):
            input_text = example.split('(TR:')[0].strip()  # Update the input text directly with the Hamshetsnag part
            break  # Only use the first clicked example

    if st.button("Tag POS"):
        with st.spinner('Processing...'):
            # Tag the input text and format output as per settings
            output = tag_pos(input_text, detailed_output)
            st.success(output)