import gradio as gr
from transformers import BlenderbotTokenizer, BlenderbotForConditionalGeneration
import torch

print("Loading MindBridge v3 model...")
model_name = "prats010/mindbridge-mental-health-model"

tokenizer = BlenderbotTokenizer.from_pretrained(model_name)
model = BlenderbotForConditionalGeneration.from_pretrained(
    model_name,
    torch_dtype=torch.float32
)
model.eval()
print("✅ Model loaded!")

def get_max_tokens(user_input):
    word_count = len(user_input.split())
    if word_count < 10:
        return 220
    else:
        return 180

def get_coping(user_input):
    lower = user_input.lower()
    if any(w in lower for w in ["anxious", "anxiety", "panic", "worry", "worried"]):
        return " Here are some ways to help: try deep breathing exercises, limit caffeine intake, practice grounding techniques like the 5-4-3-2-1 method, and consider speaking with a therapist."
    elif any(w in lower for w in ["sad", "depressed", "depression", "hopeless", "empty", "worthless"]):
        return " Some things that can help: maintain a daily routine, get sunlight and light exercise, reach out to someone you trust, and consider professional counselling."
    elif any(w in lower for w in ["sleep", "insomnia", "tired", "exhausted"]):
        return " To improve sleep: avoid screens 1 hour before bed, keep a consistent sleep schedule, try relaxation techniques like body scanning, and limit caffeine after 2pm."
    elif any(w in lower for w in ["stress", "overwhelmed", "pressure", "burnout"]):
        return " To manage stress: break tasks into smaller steps, take short breaks every 90 minutes, practice mindfulness, and talk to someone about what you are carrying."
    elif any(w in lower for w in ["lonely", "alone", "isolated", "nobody"]):
        return " To feel more connected: try joining a community or club, reach out to one person today, volunteer, or consider speaking with a counsellor who can offer consistent support."
    return ""

def chat(user_input):
    if not user_input or not user_input.strip():
        return "Hi, I'm MindBridge. How are you feeling today?"
    
    # Crisis detection
    crisis_words = ["suicide", "suicidal", "kill myself", "end my life", "want to die", "self harm", "cutting myself"]
    if any(w in user_input.lower() for w in crisis_words):
        return "I'm really concerned about what you've shared. Please reach out immediately to iCall at 9152987821 or Vandrevala Foundation at 1860-2662-345. You are not alone and help is available 24/7. [CRISIS]"
    
    inputs = tokenizer(
        user_input,
        return_tensors="pt",
        truncation=True,
        max_length=64
    )
    
    max_tokens = get_max_tokens(user_input)
    
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=max_tokens,
            num_beams=4,
            temperature=0.8,
            do_sample=True,
            top_p=0.9,
            repetition_penalty=1.3,
        )
    
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    response = response + get_coping(user_input)
    
    return response

demo = gr.Interface(
    fn=chat,
    inputs=gr.Textbox(
        label="Your message",
        placeholder="How are you feeling today?"
    ),
    outputs=gr.Textbox(label="Response"),
    title="MindBridge AI",
    description="Your mental health companion"
)

demo.launch()