import streamlit as st
from transformers import pipeline
import nltk

# Download NLTK data for sentence tokenization
nltk.download('punkt_tab')

# Load the Hugging Face pipelines
classifier = pipeline("zero-shot-classification", model="MoritzLaurer/DeBERTa-v3-base-mnli-fever-anli")
sentiment_analyzer = pipeline("sentiment-analysis", model="SarahMakk/CustomModel_amazon_sentiment_moshew_128_10k")

# Define the categories for customer feedback
CATEGORIES = ["Pricing", "Feature", "Customer Service", "Delivery", "Quality"]

# Define the fixed confidence threshold
CONFIDENCE_THRESHOLD = 0.8

# Custom CSS for background colors
st.markdown(
    """
    <style>
    /* Background color for the title */
    .title-container {
        background-color: #f0f0f0;  /* Light gray background */
        padding: 10px;
        border-radius: 5px;
    }
    /* Background color for the description (st.markdown) */
    .description-container {
        background-color: #f0f0f0;  /* Light gray background */
        padding: 10px;
        border-radius: 5px;
    }
    /* Background color for the text area (feedback_input) */
    .stTextArea textarea {
        background-color: #e6f3ff;  /* Light blue background */
        border-radius: 5px;
    }
    </style>
    """,
    unsafe_allow_html=True
)

# Streamlit app UI
# Title with icon
st.markdown(
    """
    <div class="title-container">
        <h1>📢 Customer Feedback Categorization with Sentiment Analysis</h1>
    </div>
    """,
    unsafe_allow_html=True
)

# Description with background color
st.markdown(
    """
    <div class="description-container">
        This app uses Hugging Face models to detect the topics and intent of customer feedback
        and determine the sentiment (positive👍 or negative👎) for each relevant category.
        A single feedback may belong to multiple categories, such as Pricing, Feature, and Customer Service.
    </div>
    """,
    unsafe_allow_html=True
)

# Input text box for customer feedback with background color
feedback_input = st.text_area(
    "Enter customer feedback:",
    placeholder="Type your feedback here...",
    height=200
)

# Classify button
if st.button("Classify Feedback"):
    if not feedback_input.strip():
        st.error("Please provide valid feedback text.")
    else:
        # Split the feedback into sentences
        sentences = nltk.sent_tokenize(feedback_input)
        if not sentences:
            st.error("Could not split feedback into sentences.")
            st.stop()

        # Dictionary to store results for each category
        category_results = {category: [] for category in CATEGORIES}

        # Process each sentence
        for sentence in sentences:
            # Perform zero-shot classification on the sentence
            classification_result = classifier(sentence, CATEGORIES, multi_label=True)

            # Get categories with scores above the threshold
            for label, score in zip(classification_result["labels"], classification_result["scores"]):
                if score >= CONFIDENCE_THRESHOLD:
                    # Perform sentiment analysis on the sentence
                    sentiment_result = sentiment_analyzer(sentence)
                    raw_label = sentiment_result[0]["label"]
                    sentiment_score = round(sentiment_result[0]["score"], 4)

                    # Map the raw label to NEGATIVE or POSITIVE
                    if raw_label == "LABEL_0":
                        sentiment_label = "NEGATIVE"
                        sentiment_icon = "👎"  # Thumbs-down icon for negative
                        sentiment_color = "red"  # Red color for negative
                    elif raw_label == "LABEL_1":
                        sentiment_label = "POSITIVE"
                        sentiment_icon = "👍"  # Thumbs-up icon for positive
                        sentiment_color = "green"  # Green color for positive
                    else:
                        sentiment_label = raw_label  # Fallback in case of unexpected label

                    # Store the result for the category
                    category_results[label].append({
                        "sentence": sentence,
                        "confidence": round(score, 4),
                        "sentiment": sentiment_label,
                        "sentiment_score": sentiment_score,
                        "sentiment_icon": sentiment_icon,
                        "sentiment_color": sentiment_color
                    })

        # Check if there are any relevant categories
        st.subheader("Categorized Feedback with Sentiment Analysis")
        found_categories = False

        for i, (category, results) in enumerate(category_results.items()):
            if results:  # If the category has any sentences
                found_categories = True
                st.write(f"### **{category}**")
                for result in results:
                    st.write(f"- **Sentence**: {result['sentence']}")
                    st.write(f"  - Confidence: {result['confidence']}")
                    # Use st.markdown with HTML to display the sentiment with icon and color
                    st.markdown(
                        f"  - Sentiment: {result['sentiment_icon']} "
                        f"<span style='color:{result['sentiment_color']}'>{result['sentiment']}</span> "
                        f"(Score: {result['sentiment_score']})",
                        unsafe_allow_html=True
                    )
                # Add a horizontal divider after each category (except the last one)
                if i < len(category_results) - 1 and any(category_results[cat] for cat in list(category_results.keys())[i+1:]):
                    st.markdown("---")  # Horizontal line to separate categories

        if not found_categories:
            st.warning("No categories met the confidence threshold of 0.8.")