import gradio as gr
from transformers import pipeline
import random
from datetime import datetime

# Initialize models with smaller, faster alternatives
sentiment_analyzer = pipeline(
    "sentiment-analysis",
    model="distilbert-base-uncased-finetuned-sst-2-english",
    device=-1  # Force CPU usage
)

# Pre-defined prompts and affirmations for different sentiments
PROMPT_TEMPLATES = {
    "POSITIVE": [
        "- What made this positive experience particularly meaningful to you?",
        "- How can you carry this positive energy forward?",
        "- Who would you like to share this joy with and why?"
    ],
    "NEGATIVE": [
        "- What can you learn from this challenging situation?",
        "- What small step could you take to feel better?",
        "- Who or what helps you feel supported during difficult times?"
    ]
}

AFFIRMATIONS = {
    "POSITIVE": [
        "I deserve this joy and all good things coming my way.",
        "My positive energy creates positive experiences.",
        "I choose to embrace and celebrate this moment."
    ],
    "NEGATIVE": [
        "This too shall pass, and I am growing stronger.",
        "I trust in my ability to handle challenging situations.",
        "Every experience is teaching me something valuable."
    ]
}

class JournalCompanion:
    def __init__(self):
        self.entries = []
    
    def get_prompts(self, sentiment):
        prompts = PROMPT_TEMPLATES.get(sentiment, PROMPT_TEMPLATES["POSITIVE"])
        return "\n\nReflective Prompts:\n" + "\n".join(prompts)

    def get_affirmation(self, sentiment):
        affirmations = AFFIRMATIONS.get(sentiment, AFFIRMATIONS["POSITIVE"])
        return random.choice(affirmations)

    def analyze_entry(self, entry_text):
        if not entry_text.strip():
            return ("Please write something in your journal entry.", "", "", "")

        try:
            # Perform sentiment analysis
            sentiment_result = sentiment_analyzer(entry_text)[0]
            sentiment = sentiment_result["label"].upper()
            sentiment_score = sentiment_result["score"]
        except Exception as e:
            print("Error during sentiment analysis:", e)
            return (
                "An error occurred during analysis. Please try again.",
                "Error",
                "Could not analyze sentiment due to an error.",
                "Could not generate affirmation due to an error."
            )

        entry_data = {
            "text": entry_text,
            "timestamp": datetime.now().isoformat(),
            "sentiment": sentiment,
            "sentiment_score": sentiment_score
        }
        self.entries.append(entry_data)

        # Get pre-defined responses
        prompts = self.get_prompts(sentiment)
        affirmation = self.get_affirmation(sentiment)
        sentiment_percentage = f"{sentiment_score * 100:.1f}%"
        message = f"Entry analyzed! Sentiment: {sentiment} ({sentiment_percentage} confidence)"
        
        return message, sentiment, prompts, affirmation

    def get_monthly_insights(self):
        if not self.entries:
            return "No entries yet to analyze."
            
        total_entries = len(self.entries)
        positive_entries = sum(1 for entry in self.entries if entry["sentiment"] == "POSITIVE")
        
        try:
            percentage_positive = (positive_entries / total_entries * 100)
            percentage_negative = ((total_entries - positive_entries) / total_entries * 100)
            
            insights = f"""Monthly Insights:
            Total Entries: {total_entries}
            Positive Entries: {positive_entries} ({percentage_positive:.1f}%)
            Negative Entries: {total_entries - positive_entries} ({percentage_negative:.1f}%)
            """
            return insights
        except ZeroDivisionError:
            return "No entries available for analysis."

# Rest of the code (create_journal_interface function and CSS) remains the samepad_token_id=tokenizer.eos_token_id
)
def create_journal_interface():
    journal = JournalCompanion()
    
    # Custom CSS for better styling
    custom_css = """
        /* Global styles */
        .container {
            max-width: 1200px;
            margin: 0 auto;
            padding: 20px;
        }
        
        /* Header styles */
        .header {
            text-align: center;
            margin-bottom: 2rem;
            background: linear-gradient(135deg, #6B73FF 0%, #000DFF 100%);
            padding: 2rem;
            border-radius: 15px;
            color: white;
        }
        
        /* Input area styles */
        .input-container {
            background: white;
            border-radius: 15px;
            padding: 20px;
            box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
            margin-bottom: 20px;
        }
        
        /* Output area styles */
        .output-container {
            background: #f8f9fa;
            border-radius: 15px;
            padding: 20px;
            margin-top: 20px;
        }
        
        /* Button styles */
        .custom-button {
            background: linear-gradient(135deg, #6B73FF 0%, #000DFF 100%);
            border: none;
            padding: 10px 20px;
            border-radius: 8px;
            color: white;
            font-weight: bold;
            cursor: pointer;
            transition: transform 0.2s;
        }
        
        .custom-button:hover {
            transform: translateY(-2px);
        }
        
        /* Card styles */
        .card {
            background: white;
            border-radius: 10px;
            padding: 15px;
            margin: 10px 0;
            box-shadow: 0 2px 4px rgba(0, 0, 0, 0.05);
            transition: transform 0.2s;
        }
        
        .card:hover {
            transform: translateY(-2px);
        }
        
        /* Animation for results */
        @keyframes fadeIn {
            from { opacity: 0; transform: translateY(10px); }
            to { opacity: 1; transform: translateY(0); }
        }
        
        .result-animation {
            animation: fadeIn 0.5s ease-out;
        }
        
        /* Responsive design */
        @media (max-width: 768px) {
            .container {
                padding: 10px;
            }
            
            .header {
                padding: 1rem;
            }
        }
    """
    
    with gr.Blocks(css=custom_css, title="AI Journal Companion") as interface:
        with gr.Column(elem_classes="container"):
            # Header
            with gr.Column(elem_classes="header"):
                gr.Markdown("# 📔 AI Journal Companion")
                gr.Markdown(
                    "Transform your thoughts into insights with AI-powered journaling",
                    elem_classes="subtitle"
                )
            
            # Main content
            with gr.Row():
                # Input Column
                with gr.Column(scale=1, elem_classes="input-container"):
                    entry_input = gr.Textbox(
                        label="Write Your Thoughts",
                        placeholder="Share what's on your mind...",
                        lines=8,
                        elem_classes="journal-input"
                    )
                    submit_btn = gr.Button(
                        "✨ Analyze Entry",
                        variant="primary",
                        elem_classes="custom-button"
                    )
                
                # Output Column
                with gr.Column(scale=1, elem_classes="output-container"):
                    with gr.Column(elem_classes="card result-animation"):
                        result_message = gr.Markdown(label="Analysis")
                        sentiment_output = gr.Textbox(
                            label="Emotional Tone",
                            elem_classes="sentiment-output"
                        )
                    
                    with gr.Column(elem_classes="card result-animation"):
                        prompt_output = gr.Markdown(
                            label="Reflection Prompts",
                            elem_classes="prompts-output"
                        )
                    
                    with gr.Column(elem_classes="card result-animation"):
                        affirmation_output = gr.Textbox(
                            label="Your Daily Affirmation",
                            elem_classes="affirmation-output"
                        )
            
            # Insights Section
            with gr.Row(elem_classes="insights-section"):
                with gr.Column(scale=1):
                    insights_btn = gr.Button(
                        "📊 View Monthly Insights",
                        elem_classes="custom-button"
                    )
                    insights_output = gr.Markdown(
                        elem_classes="card insights-card"
                    )

        # Event handlers
        submit_btn.click(
            fn=journal.analyze_entry,
            inputs=[entry_input],
            outputs=[
                result_message,
                sentiment_output,
                prompt_output,
                affirmation_output
            ]
        )
        
        insights_btn.click(
            fn=journal.get_monthly_insights,
            inputs=[],
            outputs=[insights_output]
        )

    return interface

if __name__ == "__main__":
    interface = create_journal_interface()
    interface.launch()