import streamlit as st
import numpy as np
import pandas as pd
import altair as alt
from transformers import BertTokenizer, BertForSequenceClassification

@st.cache_data
def get_model():
    tokenizer = BertTokenizer.from_pretrained('Dilwolf/Kakao_app-kr_sentiment')
    model = BertForSequenceClassification.from_pretrained("Dilwolf/Kakao_app-kr_sentiment")
    return tokenizer, model

tokenizer, model = get_model()

# Define the "How to Use" message
how_to_use = """
**사용하는 방법**
1. 텍스트 영역에 텍스트를 입력하세요.
2. 한국어 입력 텍스트의 예측 감정을 얻으려면 '분석' 버튼을 클릭하세요.
"""

# Functions
def main():
    st.title("BERT를 활용한 카카오톡 앱 리뷰 감성 분석")
    st.subheader("Dilshod의 포트폴리오 프로젝트")

    # Add the cover image
    st.image("img/kakaotalk.png")

    menu = ["홈", "소개"]
    choice = st.sidebar.selectbox("Menu", menu)

    # Add the "How to Use" message to the sidebar
    st.sidebar.markdown(how_to_use)

    if choice == "홈":
        st.subheader("홈")

        with st.form(key="nlpForm"):
            raw_text = st.text_area("여기에 텍스트를 입력하세요!")
            submit_button = st.form_submit_button(label="Analyze")

        # Layout
        col1, col2 = st.columns(2)
        if submit_button and raw_text:
            # Display balloons
            st.balloons()
            with col1:
                st.info("결과")
                # Tokenize the input text
                inputs = tokenizer([raw_text], padding=True, truncation=True, max_length=512, return_tensors='pt')

                # Make a forward pass through the model
                outputs = model(**inputs)

                # Get the predicted class and associated score
                predicted_class = outputs.logits.argmax().item()
                scores = outputs.logits.softmax(dim=1).detach().numpy()[0]

                # Mapping of prediction to sentiment labels
                sentiment_dict = {0: '부정적인', 1: '긍정적인'}
                sentiment_label = sentiment_dict[predicted_class]
                confidence_level = scores[predicted_class]

                # Display sentiment
                st.write(f"감정: {sentiment_label}, 신뢰 점수: {confidence_level:.2f}")

                # Emoji and sentiment image
                if predicted_class == 1:
                    st.markdown("감정 클래스: 긍정적인 :smiley:")
                    st.image("img/positive_emoji.jpg")
                else:
                    st.markdown("감정 클래스: 부정적인 :angry:")
                    st.image("img/negative_emoji.jpg")

            # Create the results DataFrame
            results_df = pd.DataFrame({
                '감정 클래스': ['부정적인', '긍정적인'],
                'Score': scores
            })

            # Create the Altair chart
            chart = alt.Chart(results_df).mark_bar(width=50).encode(
                x="감정 클래스",
                y="신뢰 점수",
                color="감정 클래스"
            )

            # Display the chart
            with col2:
                st.altair_chart(chart, use_container_width=True)
                st.write(results_df)

    else:
        st.subheader("소개")
        st.write("이것은 구글 플레이 스토어에서 가톡 모바일 앱 리뷰를 분석하기 위해 Dilshod가 개발한 감성 분석 NLP 앱입니다. 입력된 텍스트의 감정을 예측하기 위해 정밀하게 조정된 모델을 사용합니다. 이 앱은 저의 NLP 기술과 개발자 간의 협업을 보여주는 포트폴리오 프로젝트의 일부입니다.")

if __name__ == "__main__":
    main()