import gradio as gr
from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
import random
from datetime import datetime

# Initialize models
sentiment_analyzer = pipeline("sentiment-analysis", model="distilbert-base-uncased-finetuned-sst-2-english")
model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)
text_generator = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    max_new_tokens=50,
    temperature=0.7,
    top_p=0.9,
    pad_token_id=tokenizer.eos_token_id
)

class JournalCompanion:
    def __init__(self):
        self.entries = []
    
    def generate_prompts(self, sentiment):
        prompt_template = f"""Generate three reflective journal prompts for someone feeling {sentiment.lower()}. 
        Make them thoughtful and encouraging. Format them as a bullet point list."""
        
        try:
            response = text_generator(prompt_template)[0]['generated_text']
            # Extract the generated prompts after the input prompt
            prompts = response[len(prompt_template):]
            return "\n\nReflective Prompts:" + prompts
        except Exception as e:
            print("Error generating prompts:", e)
            return "\n\nReflective Prompts:\n- What thoughts and feelings are you experiencing right now?\n- How has this experience affected you?\n- What would be helpful for you at this moment?"

    def generate_affirmation(self, sentiment):
        affirmation_template = f"Generate a short, encouraging affirmation for someone feeling {sentiment.lower()}."
        
        try:
            response = text_generator(affirmation_template)[0]['generated_text']
            # Extract the generated affirmation after the input prompt
            affirmation = response[len(affirmation_template):].strip()
            return affirmation
        except Exception as e:
            print("Error generating affirmation:", e)
            return "I acknowledge my feelings and trust in my ability to handle this moment."

    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)

        # Generate responses using TinyLlama
        prompts = self.generate_prompts(sentiment)
        affirmation = self.generate_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."

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()