"""
HuggingFace Spaces that:
- loads in HanmunRoBERTa model https://huggingface.co/bdsl/HanmunRoBERTa
- optionally strips text of punctuation and unwanted charactesr
- predicts century for the input text
- Visualizes prediction scores for each century

# https://huggingface.co/blog/streamlit-spaces
# https://huggingface.co/docs/hub/en/spaces-sdks-streamlit

"""

import streamlit as st
from transformers import pipeline
from string import punctuation
import pandas as pd
import plotly.express as px
import plotly.graph_objects as go
colors = px.colors.qualitative.Plotly

# from huggingface_hub import InferenceClient
# client = InferenceClient(model="bdsl/HanmunRoBERTa")

# Load the pipeline with the HanmunRoBERTa model
model_pipeline = pipeline(task="text-classification", model="bdsl/HanmunRoBERTa")

# Streamlit app layout
title = "HanmunRoBERTa Century Classifier"
st.set_page_config(page_title=title, page_icon="📚")
st.title(title)

# Checkbox to remove punctuation
remove_punct = st.checkbox(label="Remove punctuation", value=True)

# Text area for user input
input_str = st.text_area("Input text", height=275)

# Remove punctuation if checkbox is selected
if remove_punct and input_str:
    # Specify the characters to remove
    characters_to_remove = "○□()〔〕:\"。·, ?ㆍ" + punctuation
    translating = str.maketrans('', '', characters_to_remove)
    input_str = input_str.translate(translating)

# Display the input text after processing
st.write("Processed input:", input_str)

# Predict and display the classification scores if input is provided
if st.button("Classify"):
    if input_str:
        predictions = model_pipeline(input_str)
        data = pd.DataFrame(predictions)
        data=data.sort_values(by='score', ascending=True)
        data.label = data.label.astype(str)

        
        # Displaying predictions as a bar chart
        fig = go.Figure(
            go.Bar(
                x=data.score.values,
                y=[f'{i}th Century' for i in data.label.values],
                orientation='h', 
                text=[f'{score:.3f}' for score in data['score'].values],  # Format text with 2 decimal points
                textposition='outside',  # Position the text outside the bars
                hoverinfo='text',  # Use custom text for hover info
                hovertext=[f'{i}th Century<br>Score: {score:.3f}' for i, score in zip(data['label'], data['score'])],  # Custom hover text
                marker=dict(color=[colors[i % len(colors)] for i in range(len(data))]),  # Cycle through colors

            ))
        fig.update_traces(width=0.4)

        fig.update_layout(
            height=300,  # Custom height
            xaxis_title='Score',
            yaxis_title='',
            title='Model predictions and scores',
            margin=dict(l=100, r=200, t=50, b=50),
            uniformtext_minsize=8, 
            uniformtext_mode='hide',
        )
        st.pyplot(fig=fig)
    else:
        st.write("Please enter some text to classify.")