app.py

import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig
import torch
import re

# =======================================================
# Model Loading
# =======================================================
model_name = "augtoma/qCammel-13"
print("Loading tokenizer and model...")

try:
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        device_map="auto",
        torch_dtype=torch.float16,
        trust_remote_code=True,
        low_cpu_mem_usage=True
    )
    model.eval()
    print("✅ Model loaded successfully!")
except Exception as e:
    print(f"❌ Error loading model: {e}")
    raise


# =======================================================
# Helper Functions
# =======================================================
def extract_name(text):
    text = text.lower()
    for phrase in ["yes", "i am", "i'm", "my name is", "name is", "it's", "its", "this is"]:
        text = text.replace(phrase, "")
    text = re.sub(r'[^\w\s]', '', text).strip()
    return text.title() if text and len(text) > 1 else "Patient"


def extract_age(text):
    numbers = re.findall(r'\b(\d{1,3})\b', text)
    for num in numbers:
        age = int(num)
        if 1 <= age <= 120:
            return age
    return None


def check_emergency(text):
    emergency_words = ["chest pain", "can't breathe", "suicide", "heart attack", "stroke", 
                       "severe bleeding", "unconscious", "overdose"]
    return any(word in text.lower() for word in emergency_words)


def create_initial_session():
    return {
        "name": None,
        "age": None,
        "gender": None,
        "symptoms": None,
        "duration": None,
        "medication": None,
        "other_info": [],
        "stage": "intro",
        "conversation_context": []
    }


# =======================================================
# Generate Natural Doctor Response
# =======================================================
def generate_doctor_response(session, user_message):
    """Generate a natural, conversational doctor response"""
    try:
        # Build conversational context
        context = "\n".join(session.get("conversation_context", [])[-3:])  # Last 3 exchanges
        
        prompt = f"""You are Dr. Aiden, a kind and experienced family doctor having a natural conversation with your patient.

Patient Profile:
- Name: {session['name']}
- Age: {session['age']} years old  
- Gender: {session['gender']}
- Chief Complaint: {session['symptoms']}
- Duration: {session['duration']}
- Medications: {session['medication']}

Recent conversation:
{context}

Patient: {user_message}

Dr. Aiden: """

        inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=1800).to(model.device)
        
        gen_cfg = GenerationConfig(
            temperature=0.8,
            top_p=0.92,
            max_new_tokens=300,
            repetition_penalty=1.15,
            do_sample=True,
            pad_token_id=tokenizer.pad_token_id,
            eos_token_id=tokenizer.eos_token_id
        )

        with torch.no_grad():
            output = model.generate(**inputs, generation_config=gen_cfg)

        response = tokenizer.decode(output[0], skip_special_tokens=True)
        
        # Extract just the doctor's response
        if "Dr. Aiden:" in response:
            response = response.split("Dr. Aiden:")[-1].strip()
        elif "Patient:" in response:
            response = response.split("Patient:")[0].strip()
        
        # Clean up
        response = response.strip()
        
        # If response is too short or looks like it failed, give a fallback
        if len(response) < 30 or response.startswith("You are Dr"):
            response = generate_fallback_response(session, user_message)
        
        # Add to conversation context
        session["conversation_context"].append(f"Patient: {user_message}")
        session["conversation_context"].append(f"Doctor: {response}")
        
        return response
        
    except Exception as e:
        print(f"Generation error: {e}")
        return generate_fallback_response(session, user_message)


def generate_fallback_response(session, user_message):
    """Generate a fallback response based on symptoms"""
    name = session['name']
    symptoms = session['symptoms'].lower()
    duration = session['duration']
    
    if "fever" in symptoms:
        return f"""I can see you've had this fever for {duration}, {name}. That's quite uncomfortable, I'm sure.

A fever lasting 5 days could be viral or bacterial. Let me give you some advice:

First, make sure you're drinking plenty of fluids - water, coconut water, or clear soups. Your body needs hydration when fighting infection.

For the fever itself, you can take paracetamol (like Panadol) 500mg every 6 hours if the temperature goes above 100°F. 

Get lots of rest - your body heals when you sleep. Eat light meals like khichdi, soup, or boiled vegetables.

Now here's what concerns me: if your fever goes above 103°F, you develop severe headache, difficulty breathing, or if it continues beyond 7 days, please visit a clinic immediately. These could be signs of something more serious like dengue or typhoid, which need proper testing.

How are you feeling otherwise? Any body aches, chills, or other symptoms I should know about?"""
    
    elif any(word in symptoms for word in ["cold", "cough", "flu"]):
        return f"""Okay {name}, so you've been dealing with this for {duration}. Cold and flu can really drain your energy.

Here's what I want you to do:

Stay warm and rest as much as possible. Your immune system works best when you're sleeping.

Drink warm liquids - ginger tea with honey is excellent. It soothes the throat and honey has antibacterial properties.

For congestion, try steam inhalation 2-3 times a day. Just hot water in a bowl, put a towel over your head, and breathe in the steam for 5-10 minutes.

You can take an antihistamine like Cetirizine at night if the symptoms are disturbing your sleep.

Eat vitamin C rich foods - oranges, lemons, guavas. They boost immunity.

Most colds resolve in 7-10 days. But if you develop high fever, chest pain, or difficulty breathing, that's when you need to see me in person or visit a doctor immediately.

Are you able to eat properly? Any chest discomfort?"""
    
    else:
        return f"""I understand you're concerned about {session['symptoms']}, {name}. You've been dealing with this for {duration}, which I know must be worrying.

Let me help you with this. Based on what you've told me, here are some general recommendations:

1. Keep track of your symptoms - note if they're getting better or worse
2. Stay well hydrated - aim for 8-10 glasses of water daily  
3. Get adequate rest and avoid stress
4. Eat balanced, nutritious meals

However, {name}, I want to be honest with you - some conditions really need a physical examination and possibly some tests. If your symptoms are severe, getting worse, or not improving in a few days, please visit a healthcare facility.

Is there anything specific that's worrying you the most? Or any other symptoms you haven't mentioned yet?"""


# =======================================================
# Main Response Logic
# =======================================================
def doctor_response(user_message, session):
    """Handle conversation flow naturally"""
    if not user_message or not user_message.strip():
        return "I'm here, go ahead.", session
    
    user_message_clean = user_message.strip()
    user_message_lower = user_message_clean.lower()
    
    # Emergency check
    if check_emergency(user_message_clean):
        return """⚠️ I'm concerned about what you just told me. This sounds like it could be a medical emergency.

Please call emergency services RIGHT NOW:
- Pakistan: 1122 or 115
- Or go to the nearest emergency room immediately

Don't wait. Your safety is what matters most. This chat cannot handle emergencies.""", session
    
    # Conversation flow
    try:
        if session["stage"] == "intro":
            session["stage"] = "ask_name"
            return """Hello! Come in, come in. I'm Dr. Aiden.

Please, have a seat. Before we begin, I just want you to know this is an AI consultation for guidance only. For serious conditions, you should see a doctor in person.

Now, what's your name?""", session
        
        elif session["stage"] == "ask_name":
            session["name"] = extract_name(user_message_clean)
            session["stage"] = "ask_age"
            return f"Nice to meet you, {session['name']}. And how old are you?", session
        
        elif session["stage"] == "ask_age":
            age = extract_age(user_message_clean)
            if age:
                session["age"] = age
                session["stage"] = "ask_gender"
                return f"Alright, {session['age']} years old. Are you male or female?", session
            return "Sorry, I didn't catch your age. Could you tell me in numbers? Like 25 or 30?", session
        
        elif session["stage"] == "ask_gender":
            if "male" in user_message_lower and "female" not in user_message_lower:
                session["gender"] = "male"
            elif "female" in user_message_lower:
                session["gender"] = "female"  
            else:
                return "Male or female?", session
            
            session["stage"] = "ask_symptoms"
            return f"Got it. So {session['name']}, tell me - what brings you here today? What's been bothering you?", session
        
        elif session["stage"] == "ask_symptoms":
            if len(user_message_clean) < 5:
                return "Can you describe it a bit more? I want to understand what you're going through.", session
            session["symptoms"] = user_message_clean
            session["stage"] = "ask_duration"
            return "I see. And when did this start? How long have you been feeling this way?", session
        
        elif session["stage"] == "ask_duration":
            session["duration"] = user_message_clean
            session["stage"] = "ask_medicine"
            return "Okay. Are you taking any medicines for this right now? Any tablets or syrups?", session
        
        elif session["stage"] == "ask_medicine":
            session["medication"] = user_message_clean if "no" not in user_message_lower else "No current medications"
            session["stage"] = "consult"
            
            # Natural transition to consultation
            return f"""Alright {session['name']}, let me understand this properly.

So you're {session['age']} years old, and you've been having {session['symptoms']} for about {session['duration']}. {"You're not taking anything for it right now" if "no" in user_message_lower else f"And you're currently taking {session['medication']}"}.

Let me share my thoughts with you...\n\n""" + generate_doctor_response(session, "What do you think about my condition and what should I do?"), session
        
        elif session["stage"] == "consult":
            # Natural ongoing consultation
            
            # Diet questions
            if any(word in user_message_lower for word in ["diet", "food", "eat", "meal"]):
                return get_diet_advice(session), session
            
            # Medicine questions  
            if any(word in user_message_lower for word in ["medicine", "tablet", "drug", "medication", "panadol"]):
                return get_medicine_advice(session), session
            
            # General questions
            return generate_doctor_response(session, user_message_clean), session
        
    except Exception as e:
        print(f"Error: {e}")
        return "Sorry, could you repeat that? I want to make sure I understand you correctly.", session


def get_diet_advice(session):
    """Natural diet advice"""
    name = session['name']
    symptoms = session['symptoms'].lower()
    
    if "fever" in symptoms:
        return f"""Good question about diet, {name}. When you have a fever, your body needs energy but your digestion is weak.

Here's what I recommend:

**For breakfast**: Keep it light - maybe some oatmeal porridge or a boiled egg with a banana. Easy to digest.

**Lunch and dinner**: Khichdi is excellent - rice and dal together. Or you can have plain rice with some boiled vegetables or light chicken soup.

**Throughout the day**: This is very important - drink lots of fluids. Water, coconut water, lemon water. Your body loses water when you have fever.

**Avoid these**: No fried foods, no spicy food, no heavy curries. Your stomach can't handle them right now.

Little meals, multiple times a day. Don't force yourself to eat if you're not hungry, but definitely keep drinking fluids.

Does that make sense? Any other questions about food?"""
    
    elif "cold" in symptoms or "cough" in symptoms:
        return f"""Ah yes, diet is important for recovery, {name}.

When you have cold and cough:

**Hot is good**: Warm soups, especially chicken soup - it's not just comfort food, it actually helps with congestion. Ginger tea with honey is excellent.

**Fruits**: Oranges, kinnow, any citrus fruits. They have vitamin C which helps your immune system fight the infection.

**Regular meals**: Dal rice, roti sabzi - normal food is fine. Just make sure it's warm, not cold from the fridge.

**Avoid**: Cold drinks, ice cream, banana (can increase mucus), dairy products if you have lot of phlegm.

**Drink**: Warm water throughout the day. You can add a bit of honey and lemon to it.

And one home remedy - at night before sleeping, have a glass of warm milk with half teaspoon of turmeric. It helps with throat irritation.

Anything else you'd like to know?"""
    
    else:
        return f"""For what you're experiencing, {name}, I'd say maintain a balanced diet.

Eat fresh vegetables, fruits, whole grains. Include protein - dal, chicken, fish, eggs. Your body needs good nutrition to heal.

Stay hydrated - very important. At least 8-10 glasses of water daily.

Avoid processed foods, too much sugar, fried items.

Eat at regular times, don't skip meals.

Is there something specific about your diet you're concerned about?"""


def get_medicine_advice(session):
    """Natural medicine advice"""
    name = session['name']
    symptoms = session['symptoms'].lower()
    
    response = f"""Alright {name}, regarding medications:\n\n"""
    
    if "fever" in symptoms:
        response += """For your fever, you can take **Paracetamol 500mg** (Panadol is the common brand here).

Take one tablet every 6 hours, but only when your temperature is above 100°F. Don't take it more than 4 times in 24 hours.

Take it after eating something, even if just a few biscuits. Don't take it on empty stomach.

If the fever is not coming down with Paracetamol, or if it's going very high (above 103°F), that's a red flag. You need to see a doctor in person then."""
    
    elif "cold" in symptoms or "cough" in symptoms:
        response += """For cold and cough:

**For runny nose/sneezing**: Cetirizine 10mg, one tablet at night before sleeping. It'll help you rest better.

**For cough**: If it's dry cough - any cough syrup with dextromethorphan. If you have phlegm - syrup with guaifenesin (expectorant). Take as per bottle instructions.

**For blocked nose**: You can use saline nasal drops.

These are all over-the-counter, you can get them from any pharmacy.

But listen carefully - if you develop fever, chest pain, or difficulty breathing, stop these and consult a doctor immediately."""
    
    else:
        response += """I can give you general guidance, but I'd be more comfortable if you saw a doctor in person for proper examination and prescription.

Some conditions need specific tests or examinations before recommending medicines.

For now, stay comfortable, rest well, and monitor your symptoms. If anything worsens, please visit a clinic."""
    
    response += "\n\nAny other questions for me?"
    return response


# =======================================================
# Gradio Interface
# =======================================================
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.HTML("""
    <div style="text-align:center; background: linear-gradient(135deg, #2193b0 0%, #6dd5ed 100%); color:white; padding:30px; border-radius:15px;">
        <h1 style="margin:0; font-size: 2.5em;">🩺 Dr. Aiden's Clinic</h1>
        <p style="margin:10px 0 0 0; font-size: 1.2em;">Your friendly AI doctor - here to listen and help</p>
    </div>
    """)
    
    session_state = gr.State(create_initial_session)
    
    chatbot = gr.Chatbot(
        label="💬 Consultation Room",
        height=550,
        type='messages',
        avatar_images=(
            "https://cdn-icons-png.flaticon.com/512/706/706830.png",
            "https://cdn-icons-png.flaticon.com/512/3774/3774299.png"
        )
    )
    
    with gr.Row():
        user_input = gr.Textbox(
            placeholder="Type your message here...", 
            label="", 
            lines=2,
            scale=4
        )
    
    with gr.Row():
        send_btn = gr.Button("💬 Send", variant="primary", scale=1)
        clear_btn = gr.Button("🔄 New Patient", variant="secondary", scale=1)
    
    gr.HTML("""
    <div style="margin-top: 15px; padding: 12px; background-color: #fff3cd; border-radius: 8px; font-size: 0.9em;">
        <strong>⚠️ Important:</strong> This is an AI simulation. Always consult a real doctor for medical conditions. 
        For emergencies call 1122 (Pakistan) or 911.
    </div>
    """)
    
    def respond(message, history, session):
        if history is None:
            history = []
        if not message or not message.strip():
            return "", history, session
        
        response, updated_session = doctor_response(message, session)
        
        history.append({"role": "user", "content": message})
        history.append({"role": "assistant", "content": response})
        
        return "", history, updated_session
    
    def reset():
        return [], create_initial_session()
    
    send_btn.click(respond, [user_input, chatbot, session_state], [user_input, chatbot, session_state])
    user_input.submit(respond, [user_input, chatbot, session_state], [user_input, chatbot, session_state])
    clear_btn.click(reset, None, [chatbot, session_state], queue=False)

if __name__ == "__main__":
    print("🩺 Dr. Aiden's Clinic is now open...")
    demo.queue()
    demo.launch(share=True)