# V3
# Improve speed and user experience 

import streamlit as st
from transformers import pipeline
from langdetect import detect

# 加载翻译 pipeline
@st.cache_resource
def load_translation_pipeline():
    return pipeline("translation", model="facebook/m2m100_418M", max_length=256)

# 加载情感分析 pipeline
@st.cache_resource
def load_sentiment_pipeline():
    return pipeline("sentiment-analysis", model="Rocky080808/finetuned-distilbert-base-uncased-finetuned-sst-2-english", max_length=256)

# 定义语言映射
language_name_map = {
    'en': "English",
    'zh-cn': "Chinese (Simplified)",
    'zh-tw': "Chinese (Traditional)",
    'ja': "Japanese",
    'de': "German",
    'es': "Spanish",
    'fr': "French"
}

# 翻译到英语的函数
def translate_to_english(text, translation_pipeline):
    detected_language = detect(text)
    
    # 语言映射
    language_map = {
        'en': "en",     # 英语直接通过
        'zh-cn': "zh",  # Simplified Chinese
        'zh-tw': "zh",  # Traditional Chinese
        'ja': "ja",     # Japanese
        'de': "de",     # German
        'es': "es",     # Spanish
        'fr': "fr"      # French
    }

    if detected_language not in language_map:
        return None, "Unsupported language"
    
    # 如果检测到是英语，直接返回原文本
    if detected_language == 'en':
        return text, "en"
    
    # 翻译为英语
    translated_text = translation_pipeline(text, src_lang=language_map[detected_language], tgt_lang="en")
    return translated_text[0]['translation_text'], language_name_map.get(detected_language, detected_language)

# 主程序逻辑
def main():
    # 加载翻译和情感分析模型
    translation_pipeline = load_translation_pipeline()
    sentiment_pipeline = load_sentiment_pipeline()

    st.title("Global Customer Reviews Sentiment Analyzer")
    st.write("Analyze customer sentiment by their reviews. Please input the customer reviews to get the sentiment analysis result.")
    st.write("Support 6 languages: English, Chinese, Japanese, German, Spanish and French.")
    st.write("For example, by inputting: I like the product very much!")
    st.write("The application will tell you: Very satisfied, the customer is very likely to return and recommend.")

    user_input = st.text_input("Enter customer reviews here and press Analyze:")

    # 用户点击分析按钮后触发
    if st.button("Analyze"):
        if user_input:
            # 翻译或直接处理英语
            translated_text, detected_language = translate_to_english(user_input, translation_pipeline)
            
            if detected_language == "Unsupported language":
                st.write("The input language is not supported. Please use Chinese, Japanese, German, Spanish, or French.")
            else:
                # 显示检测语言和翻译结果（如果需要翻译）
                st.write(f"Detected language: {language_name_map.get(detected_language, detected_language)}")
                st.write(f"Translated Text: {translated_text}" if detected_language != "English" else f"Original Text: {translated_text}")

                # 情感分析
                result = sentiment_pipeline(translated_text)
                label_str = result[0]["label"]
                label = int(label_str.split("_")[-1])
                confidence = result[0]["score"]

                # 情感结果映射
                label_to_text = {
                    0: "Very dissatisfied, immediate follow-up is required.",
                    1: "Dissatisfied, please arrange follow-up.",
                    2: "Neutral sentiment, further case analysis is needed.",
                    3: "Satisfied, the customer may return for a purchase.",
                    4: "Very satisfied, the customer is very likely to return and recommend."
                }
                sentiment_text = label_to_text.get(label, "Unrecognized sentiment")

                st.write(f"Sentiment Analysis Result: {sentiment_text}")
                # st.write(f"Confidence Score: {confidence:.2f}")

if __name__ == "__main__":
    main()