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
        output = []
        for label in sentence.get_labels():
            output.append(
                f"{label.data_point.text}: {label.value} ({label.score:.2f})"
            )
        return "\n".join(output)
    else:
        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 text area
    input_text = st.text_area("Enter a text:", height=100, value=st.session_state.get('input_text', 'Put example text here.'))

    # Tag POS button with unique color styling
    if st.button("Tag POS", key="tag_pos"):
        with st.spinner('Processing...'):
            output = tag_pos(input_text, detailed_output)
            st.success(output)

    # Example Sentences and Translations
    example_sentences = [
        ("tuute acertsetser topoldetser aaav ta.", "Kâğıdı büzüştürdün attın. Oldu mu?"),
        ("Baran u Baden teran.", "Baran ve Bade koştu."),
        ("Onurun ennush nu İremin terchushe intzi shad kızdırmısh aaav.", "Onur'un düşüşü ve İrem'in koşuşu beni kızdırdı."),
    ]

    st.write("## Example Sentences:")

    for hamshetsnag, turkish in example_sentences:
        if st.button(f"Use: {hamshetsnag}", key=hamshetsnag):
            st.session_state['input_text'] = hamshetsnag  # Update input text
            break

        st.markdown(f'<p style="font-size:12px;">(TR: {turkish})</p>', unsafe_allow_html=True)  # Display translation in smaller font