# 1. Import necessary libraries
import gradio as gr
from transformers import pipeline

# 2. Load the NEW, more specialized AI model
emotion_classifier = pipeline(
    "text-classification",
    model="mental/mental-roberta-base", # <-- UPGRADED MODEL
    top_k=None
)

# 3. Define the function to process the input and return outputs
def analyze_text(text_input):
    """
    This function takes text, passes it to the AI model,
    and then formats the results for a better user experience.
    """
    # Get the raw predictions from the model
    predictions = emotion_classifier(text_input)[0]

    # Create a dictionary to hold the scores for key indicators
    key_indicators = {
        'sadness': 0, 'anger': 0, 'fear': 0, 'joy': 0
    }
    for emotion in predictions:
        if emotion['label'] in key_indicators:
            key_indicators[emotion['label']] = round(emotion['score'], 3)
    
    # --- NEW: Interpretation Logic ---
    # Find the dominant emotion among our key indicators
    if not key_indicators:
        dominant_emotion = "neutral"
    else:
        dominant_emotion = max(key_indicators, key=key_indicators.get)

    interpretation_text = ""
    resource_links = ""

    if key_indicators.get(dominant_emotion, 0) > 0.4: # Set a threshold
        if dominant_emotion in ['sadness', 'fear', 'anger']:
            interpretation_text = (
                f"The analysis shows a strong presence of **{dominant_emotion}**. "
                "These feelings can be challenging. Remember, it's okay to seek support."
            )
            # --- NEW: Provide helpful resources ---
            resource_links = """
            **Please consider reaching out to a professional:**
            - **Vandrevala Foundation:** [vandrev ফাউন্ডেশন](https://www.vandrevalafoundation.com/) (India)
            - **NIMHANS Centre for Well-Being:** [NIMHANS](http://www.nimhans.ac.in/well-being-centre/) (Bengaluru)
            - **Find a Helpline:** [findahelpline.com](https://findahelpline.com/) (Global)
            """
        elif dominant_emotion == 'joy':
            interpretation_text = (
                f"The text shows strong indicators of **joy**. It's wonderful to see such positive expression."
            )

    # Return the scores, the interpretation, and the resources
    return key_indicators, interpretation_text, resource_links

# 4. Create the Gradio web interface with new components
app_interface = gr.Interface(
    fn=analyze_text,
    inputs=gr.Textbox(
        lines=8,
        label="Social Media Post",
        placeholder="Type or paste a social media post here..."
    ),
    # --- NEW: Multiple output components for a richer display ---
    outputs=[
        gr.Label(label="Key Emotional Indicators"),
        gr.Markdown(label="Interpretation"),
        gr.Markdown(label="Resources")
    ],
    title="Enhanced Student Wellness Analyzer 🧠✨",
    description="""
    This upgraded tool uses a specialized AI to analyze text for emotional indicators.
    **Disclaimer:** This is **not a medical diagnostic tool**. It is an AI demonstration.
    If you are struggling, please seek help from a qualified professional.
    """,
    examples=[
        ["I'm so behind on all my assignments and the exams are next week. I don't know how I'm going to manage all this pressure."],
        ["Feeling completely isolated and lonely on campus. It seems like everyone else has their friend group figured out."],
        ["Finished my project and I'm so proud of how it turned out! The hard work really paid off."]
    ]
)

# 5. Launch the app
app_interface.launch()