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.
        The feedback is split into sentences, and each sentence is categorized and analyzed for sentiment.
        Only categories with a confidence score >= 0.8 are displayed.
    </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,
    value="I was shocked to see the price tag on this new gadget—it’s way too expensive for what it offers, especially compared to competitors! Despite the issues I faced with my order, the customer service team's effort to rectify the situation was commendable, though their follow-up could have used some improvement for full satisfaction."
)

# 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
                    else raw_label == "LABEL_1":
                        sentiment_label = "POSITIVE"
                        sentiment_icon = "👍"  # Thumbs-up icon for positive
                        sentiment_color = "green"  # Green color for positive

                    # 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.")