# app.py
import streamlit as st
from groq import Groq
from textblob import TextBlob
from transformers import pipeline
import re
import random
from datetime import datetime
from reportlab.lib.pagesizes import letter
from reportlab.pdfgen import canvas
import io

# Initialize components with proper error handling
try:
    groq_client = Groq(api_key=st.secrets["GROQ_API_KEY"])
except KeyError:
    st.error("GROQ_API_KEY missing in secrets! Add it in Hugging Face settings.")
    st.stop()

# Initialize psychological models using verified public models
try:
    # Personality insights using zero-shot classification
    personality_analyzer = pipeline(
        "zero-shot-classification",
        model="facebook/bart-large-mnli"
    )
    
    # Emotion detection model
    emotion_classifier = pipeline(
        "text-classification", 
        model="j-hartmann/emotion-english-distilroberta-base"
    )
except Exception as e:
    st.error(f"Model loading error: {str(e)}")
    st.stop()

# Configure Streamlit
st.set_page_config(page_title="🧠 Mind Mapper", layout="wide", page_icon="🤖")

# Custom CSS
st.markdown("""
<style>
@keyframes rainbow {
  0% { color: #ff0000; }
  20% { color: #ff8000; }
  40% { color: #ffff00; }
  60% { color: #00ff00; }
  80% { color: #0000ff; }
  100% { color: #ff00ff; }
}
.personality-title {
  animation: rainbow 3s infinite;
  font-size: 2.5em !important;
}
.progress-bar {
  height: 25px;
  border-radius: 15px;
  background: linear-gradient(90deg, #FF6B6B 0%, #4ECDC4 100%);
}
.tip-box {
  border-left: 4px solid #4CAF50;
  padding: 10px 15px;
  margin: 10px 0;
  background: #f8f9fa;
  border-radius: 5px;
}
</style>
""", unsafe_allow_html=True)

# Personality traits and questions
OCEAN_TRAITS = {
    "openness": ["curious", "creative", "adventurous"],
    "conscientiousness": ["organized", "responsible", "disciplined"],
    "extraversion": ["outgoing", "energetic", "sociable"],
    "agreeableness": ["friendly", "compassionate", "cooperative"],
    "neuroticism": ["sensitive", "nervous", "emotional"]
}

QUESTION_BANK = [
    {"text": "How do you typically approach new experiences? 🌍", "trait": "openness"},
    {"text": "Describe your ideal weekend vs reality 🏝️", "trait": "conscientiousness"},
    {"text": "How do you recharge after social events? 🪫", "trait": "extraversion"},
    {"text": "Your reaction to someone disagreeing with you? 💢", "trait": "agreeableness"},
    {"text": "How do you handle unexpected changes? 🌀", "trait": "neuroticism"}
]

def analyze_psychology(text):
    """Analyze personality traits using zero-shot classification"""
    candidate_labels = list(OCEAN_TRAITS.keys())
    result = personality_analyzer(text, candidate_labels, multi_label=True)
    return {label: score for label, score in zip(result['labels'], result['scores'])}

def analyze_emotions(text):
    """Detect emotional tone"""
    return emotion_classifier(text[:512])

def generate_quote(traits):
    """Generate personality-specific quote"""
    prompt = f"""Create a motivational quote for someone with these traits:
{traits}
Make it inspirational and include an emoji."""
    
    response = groq_client.chat.completions.create(
        model="mixtral-8x7b-32768",
        messages=[{"role": "user", "content": prompt}],
        temperature=0.7
    )
    return response.choices[0].message.content

def generate_tips(traits):
    """Generate evidence-based psychological tips for all traits"""
    tips = []
    
    # Openness
    if traits.get('openness', 0) > 0.7:
        tips.extend([
            "🎨 Creative boost: Try 'morning pages' - write 3 stream-of-consciousness pages each morning",
            "🌍 Novelty diet: Schedule one new experience weekly (new route, food, or activity)"
        ])
    else:
        tips.append("📚 Structured exploration: Try a 'learning hour' weekly on a fixed schedule about random topics")

    # Conscientiousness
    if traits.get('conscientiousness', 0) > 0.6:
        tips.extend([
            "⏱️ Time boxing: Use the 52/17 rule (52 mins work, 17 mins break) for optimal productivity",
            "📦 Declutter challenge: Remove 10 unnecessary items from your workspace daily"
        ])
    else:
        tips.extend([
            "🎯 Micro-goals: Break tasks into 5-minute chunks with immediate rewards",
            "📅 Habit stacking: Attach new habits to existing routines (e.g., meditate after brushing teeth)"
        ])

    # Extraversion
    if traits.get('extraversion', 0) > 0.6:
        tips.extend([
            "👥 Energy sharing: Schedule weekly 'idea exchanges' with friends",
            "🎤 Social charging: Host monthly themed dinners to fuel your social battery"
        ])
    else:
        tips.extend([
            "📔 Reflection ritual: Keep an 'insight journal' for deep thinking sessions",
            "🧠 Solo innovation: Try monthly 'solo design sprints' for personal projects"
        ])

    # Agreeableness
    if traits.get('agreeableness', 0) > 0.6:
        tips.extend([
            "🛡️ Boundary practice: Start sentences with 'I choose to...' instead of 'I have to...'",
            "💪 Diplomatic assertiveness: Use the 'DESC' script (Describe, Express, Specify, Consequences)"
        ])
    else:
        tips.extend([
            "🤝 Perspective walks: Have weekly conversations where you only ask questions",
            "❤️ Compassion training: Practice daily '3-good-things' gratitude journaling"
        ])

    # Neuroticism
    if traits.get('neuroticism', 0) > 0.6:
        tips.extend([
            "🛑 Anxiety containment: Create a 'worry window' (15 mins/day dedicated to problem-solving)",
            "🌊 Emotional surfing: Practice observing feelings like waves (notice-name-allow technique)"
        ])
    else:
        tips.extend([
            "🎭 Vulnerability practice: Share one authentic feeling daily with someone",
            "🚧 Comfort zone challenges: Do something mildly uncomfortable weekly"
        ])

    # General well-being tips
    tips.extend([
        "💤 Sleep hygiene: Follow 10-3-2-1-0 formula (10h before bed: caffeine, 3h: food, 2h: work, 1h: screens, 0 snooze)",
        "🧠 Cognitive refueling: Practice 20-20-20 rule (every 20 mins, look 20 ft away for 20 sec)",
        "❤️ Self-compassion break: When stressed, say 'This is hard, but I'm doing my best'"
    ])

    return tips[:15]  # Return top 15 most relevant tips

def create_pdf_report(report):
    """Generate downloadable PDF report"""
    buffer = io.BytesIO()
    p = canvas.Canvas(buffer, pagesize=letter)
    
    p.setFont("Helvetica-Bold", 16)
    p.drawString(100, 750, "🌟 Psychological Profile Report 🌟")
    p.setFont("Helvetica", 12)
    
    y_position = 700
    content = [
        f"Generated on: {datetime.now().strftime('%Y-%m-%d %H:%M')}",
        "",
        "Personality Traits:",
        *[f"- {trait.upper()}: {score:.2f}" for trait, score in report['traits'].items()],
        "",
        "Emotional Analysis:",
        *[f"- {emo['label']}: {emo['score']:.2f}" for emo in report['emotions']],
        "",
        "Personalized Quote:",
        report['quote'],
        "",
        "Growth Tips:",
        *[f"{i+1}. {tip}" for i, tip in enumerate(report['tips'])]
    ]
    
    for line in content:
        p.drawString(100, y_position, line)
        y_position -= 15
    
    p.save()
    buffer.seek(0)
    return buffer

# Session state management
if 'responses' not in st.session_state:
    st.session_state.responses = []
if 'current_q' not in st.session_state:
    st.session_state.current_q = 0

# Main UI
st.title("🧠 Mind Mapper")
st.markdown("### Your Personal Psychology Assistant 🔍✨")

# Progress tracker
progress = st.session_state.current_q / len(QUESTION_BANK)
st.markdown(f"""
<div style="background: #f0f2f6; border-radius: 15px; padding: 5px;">
    <div class="progress-bar" style="width: {progress*100}%; height: 25px;"></div>
</div>
""", unsafe_allow_html=True)

# Question flow
if st.session_state.current_q < len(QUESTION_BANK):
    q = QUESTION_BANK[st.session_state.current_q]
    
    with st.chat_message("assistant"):
        st.markdown(f"### {q['text']}")
        user_input = st.text_input("Your response:", key=f"q{st.session_state.current_q}")
        
        if st.button("Next ➡️"):
            st.session_state.responses.append(user_input)
            st.session_state.current_q += 1
            st.rerun()
else:
    # Generate report
    combined_text = "\n".join(st.session_state.responses)
    traits = analyze_psychology(combined_text)
    emotions = analyze_emotions(combined_text)
    
    report = {
        "traits": traits,
        "emotions": emotions,
        "quote": generate_quote(traits),
        "tips": generate_tips(traits)
    }
    
    st.balloons()
    st.markdown(f"## <div class='personality-title'>🌟 Your Psychological Profile 🌟</div>", unsafe_allow_html=True)
    
    # Trait Visualization
    with st.expander("🧩 Personality Breakdown"):
        cols = st.columns(5)
        for i, (trait, score) in enumerate(report['traits'].items()):
            cols[i].metric(label=trait.upper(), value=f"{score:.0%}")
    
    # Emotional Analysis
    with st.expander("🎭 Emotional Landscape"):
        emotion = max(report['emotions'], key=lambda x: x['score'])
        st.markdown(f"**Dominant Emotion**: {emotion['label']} ({emotion['score']:.0%})")
        st.progress(emotion['score'])
    
    # Recommendations
    col1, col2 = st.columns(2)
    with col1:
        st.markdown("### 💬 Personalized Wisdom")
        st.success(f'"{report["quote"]}"')
    with col2:
        st.markdown("### 🛠️ Actionable Psychology Hacks")
        for i, tip in enumerate(report['tips']):
            st.markdown(f"""
            <div class="tip-box">
                <strong>#{i+1}</strong> {tip}
            </div>
            """, unsafe_allow_html=True)
    
    # PDF Report
    pdf_buffer = create_pdf_report(report)
    st.download_button(
        "📥 Download Full Report", 
        data=pdf_buffer, 
        file_name="psych_profile.pdf",
        mime="application/pdf"
    )

# Sidebar
with st.sidebar:
    st.markdown("## 🌈 How It Works")
    st.markdown("""
    1. Answer 5 psychology-based questions
    2. Get instant personality analysis
    3. Receive emotional insights
    4. Download personalized report
    
    **Science Behind It:**
    - Zero-shot personality classification
    - Emotion recognition (RoBERTa)
    - Cognitive behavioral therapy principles
    """)