import gradio as gr
import random
from huggingface_hub import InferenceClient
from sentence_transformers import SentenceTransformer
import torch

client = InferenceClient('Qwen/Qwen2.5-72B-Instruct')

# Step 1
with open("Mental health chatbot text.txt", "r", encoding="utf-8") as f:
    mental_health_text = f.read()

# Step 2: Preprocess text into sentence chunks
def preprocess_text(text):
    cleaned_text = text.strip()
    sentences = [s.strip() for s in cleaned_text.split('.') if s.strip()]
    sentence_chunks = [s.strip() for s in sentences if len(s.strip()) > 10]
    
    combined_chunks = []
    for i in range(0, len(sentences), 2):
        chunk = '. '.join(sentences[i:i+3]).strip()
        if len(chunk) > 20:
            combined_chunks.append(chunk)
    
    paragraphs = [p.strip() for p in cleaned_text.split('\n\n') if p.strip()]
    paragraph_chunks = [p for p in paragraphs if len(p) > 30]
    
    all_chunks = sentence_chunks + combined_chunks + paragraph_chunks
    
    seen = set()
    final_chunks = []
    for chunk in all_chunks:
        if chunk not in seen and len(chunk) > 15:
            seen.add(chunk)
            final_chunks.append(chunk)
    
    print(f"Created {len(final_chunks)} chunks using advanced strategy")
    print(f"Sample chunks: {final_chunks[:3]}")
    return final_chunks

cleaned_chunks = preprocess_text(mental_health_text)

# Step 3: Convert chunks into embeddings
model = SentenceTransformer('all-MiniLM-L6-v2')

def create_embeddings(text_chunks):
    chunk_embeddings = model.encode(text_chunks, convert_to_tensor=True)
    print(f"Embeddings shape: {chunk_embeddings.shape}")
    return chunk_embeddings

chunk_embeddings = create_embeddings(cleaned_chunks)

# Step 4: Retrieve top matching chunks
def get_top_chunks(query, chunk_embeddings, text_chunks, top_k=3):
    query_embedding = model.encode(query, convert_to_tensor=True)
    query_norm = query_embedding / query_embedding.norm()
    chunks_norm = chunk_embeddings / chunk_embeddings.norm(dim=1, keepdim=True)
    similarities = torch.matmul(chunks_norm, query_norm)
    top_indices = torch.topk(similarities, k=top_k).indices
    return [text_chunks[i] for i in top_indices]

# Step 5: Relevance checker
def is_mental_health_related(query):
    mental_health_keywords = [
        'anxiety', 'depression', 'stress', 'mental health', 'therapy', 'counseling',
        'mood', 'emotions', 'feelings', 'wellbeing', 'self-care', 'mindfulness',
        'meditation', 'coping', 'support', 'psychology', 'psychiatry', 'bipolar',
        'trauma', 'PTSD', 'panic', 'worry', 'sad', 'happy', 'angry', 'fear',
        'self-esteem', 'confidence', 'resilience', 'healing', 'recovery',
        'mental', 'emotional', 'psychological', 'behavioral', 'cognitive'
    ]
    
    query_lower = query.lower()
    return any(keyword in query_lower for keyword in mental_health_keywords)

queries = [
    "Managing daily stress and anxiety",
    "Building healthy coping mechanisms", 
    "Practicing mindfulness and self-care"
]

for q in queries:
    print(f"\nQuery: {q}")
    results = get_top_chunks(q, chunk_embeddings, cleaned_chunks)
    for idx, res in enumerate(results, 1):
        print(f"Result {idx}: {res}")

def respond(message, history):
    try:
        top_results = get_top_chunks(message, chunk_embeddings, cleaned_chunks)
        print(top_results)
        
        messages = [{"role": "system", "content": f"You are a friendly chatbot. You give people advice about mental health. Base your response on the following information: {top_results}"}]
        
        if history:
            for user_msg, bot_msg in history:
                messages.append({"role": "user", "content": user_msg})
                messages.append({"role": "assistant", "content": bot_msg})
        
        messages.append({"role": "user", "content": message})
        
        response = client.chat_completion(messages)
        
        if hasattr(response, 'choices') and response.choices:
            return response.choices[0].message.content.strip()
        elif isinstance(response, dict) and 'choices' in response:
            return response['choices'][0]['message']['content'].strip()
        else:
            return str(response).strip()
            
    except Exception as e:
        print(f"Error: {e}")
        return "I'm having trouble right now. Please try again."

def start_breathing_game():
    return (
        """
        <style>
        .breathing-circle {
            width: 120px;
            height: 120px;
            border-radius: 50%;
            display: flex;
            align-items: center;
            justify-content: center;
            margin: 30px auto;
            font-weight: bold;
            font-size: 14px;
            animation: breathing 8s ease-in-out infinite;
        }

        @keyframes breathing {
            0%, 100% {
                transform: scale(1);
                background: linear-gradient(135deg, #a7f3d0 0%, #d1fae5 100%);
                color: #065f46;
                box-shadow: 0 4px 20px rgba(167, 243, 208, 0.4);
                content: "🍃 Exhale...";
            }
            50% {
                transform: scale(1.8);
                background: linear-gradient(135deg, #34d399 0%, #10b981 100%);
                color: white;
                box-shadow: 0 12px 40px rgba(52, 211, 153, 0.5);
                content: "🌸 Inhale...";
            }
        }

        .breathing-container {
            background: linear-gradient(135deg, rgba(255,255,255,0.9) 0%, rgba(240,253,244,0.9) 100%);
            border-radius: 20px;
            padding: 40px 30px;
            box-shadow: 0 10px 25px -5px rgba(0, 0, 0, 0.1);
            backdrop-filter: blur(10px);
            min-height: 400px;
            display: flex;
            flex-direction: column;
            justify-content: center;
        }
        </style>
        
        <div class="breathing-container">
            <div style="text-align: center;">
                <div class="breathing-circle">Breathe</div>
                <p style="color: #065f46; font-weight: 600; margin: 25px 0 15px 0;"><strong>Follow the gentle rhythm:</strong></p>
                <p style="color: #059669; margin: 8px 0;">🌸 When it grows → Inhale slowly and deeply</p>
                <p style="color: #059669; margin: 8px 0;">🍃 When it shrinks → Exhale and release tension</p>
                <p style="color: #047857; font-style: italic; margin-top: 20px;"><em>Take your time and focus on your breath. You're safe here.</em></p>
            </div>
        </div>
        """,
        "🌟 **Perfect! Take a moment for yourself.** Follow the breathing circle above - it will guide you through slow, calming breaths. Let the gentle rhythm help you find your center. Take as long as you need."
    )

# Theme
soothing_theme = gr.themes.Soft(
    primary_hue="emerald",
    secondary_hue="green",
    neutral_hue="slate",
    font=gr.themes.GoogleFont("Inter")
)

custom_css = """
.gradio-container {
    background: linear-gradient(135deg, #f0fdf4 0%, #ecfdf5 50%, #f9fafb 100%) !important;
}
"""

with gr.Blocks(theme=soothing_theme, css=custom_css, title="Otium") as demo:
    gr.HTML("""
    <div style="text-align: center; padding: 20px;">
        <h1 style="color: #065f46;">Otium</h1>
        <p style="color: #059669;">Your gentle companion for mental wellness support</p>
    </div>
    """)
    
    with gr.Tabs():
        with gr.TabItem("💬 Chat"):
            chatbot = gr.ChatInterface(
                respond,
                chatbot=gr.Chatbot(
                    height=500,
                    show_label=False,
                    container=True,
                    bubble_full_width=False,
                    avatar_images=("🧘‍♀️", "🌱")
                ),
                textbox=gr.Textbox(
                    placeholder="Share what's on your mind... 💭",
                    container=False,
                    scale=7
                ),
                submit_btn="Send 🌸"
            )
            
        with gr.TabItem("🫁 Breathe Break"):
            gr.Markdown("""
            ### 🌿 Mindful Breathing Exercise
            *Take a moment to center yourself with this gentle breathing practice.*
            """)
            
            start_button = gr.Button(
                "🌸 Start Breathing Exercise", 
                variant="primary",
                size="lg"
            )
            game_display = gr.HTML()
            game_text = gr.Markdown()
            
            start_button.click(
                fn=start_breathing_game,
                outputs=[game_display, game_text]
            )

demo.launch(share=False, inbrowser=True)