app.py

import gradio as gr
import requests
import os
from datetime import datetime
from reportlab.lib.pagesizes import letter
from reportlab.pdfgen import canvas
from transformers import pipeline
import torch

# === Groq API Setup ===
GROQ_API_KEY = os.getenv("GROQ_API_KEY")
if not GROQ_API_KEY:
    raise ValueError("GROQ_API_KEY environment variable is not set.")
GROQ_MODEL = "llama3-8b-8192"

# === Medical Classifier ===
# Removed classifier usage, so no filtering/warnings.

def doctor_twin_light(prompt, category):
    # Removed the medical classifier check here

    system_prompt = (
        "You are Doctor Twin, a virtual AI health assistant. "
        "You provide friendly, general health and wellness advice. "
        "Also check if question is not medical related, do give warning that it is not relevant to medical."
        "You never diagnose or prescribe. Always include a disclaimer to consult a real doctor."
    )

    user_prompt = f"Category: {category}\nPatient: {prompt}\nAdvice:"

    headers = {
        "Authorization": f"Bearer {GROQ_API_KEY}",
        "Content-Type": "application/json"
    }

    data = {
        "model": GROQ_MODEL,
        "messages": [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt}
        ],
        "temperature": 0.7,
        "max_tokens": 256
    }

    try:
        response = requests.post("https://api.groq.com/openai/v1/chat/completions", headers=headers, json=data)
        response.raise_for_status()
        reply = response.json()["choices"][0]["message"]["content"]
        return f"{reply}\n\n⚠️ Disclaimer: AI-generated advice. Always consult a licensed doctor."
    except Exception as e:
        return f"❌ Error: {str(e)}"

# === OTC Prescription Generator ===
def generate_otc_prescription(name, symptoms):
    date = datetime.now().strftime('%Y-%m-%d')
    content = f"""📄 Prescription Note
Patient: {name}
Date: {date}
Symptoms: {symptoms}
Suggested OTC Medicine: Cetirizine 10mg (once at night)
Instructions: Take after food. Stay hydrated.
Caution: This is an AI-generated suggestion. Please consult a licensed doctor if symptoms persist."""

    filename = f"prescription_{datetime.now().strftime('%Y%m%d%H%M%S')}.pdf"
    filepath = os.path.join("prescriptions", filename)
    os.makedirs("prescriptions", exist_ok=True)

    c = canvas.Canvas(filepath, pagesize=letter)
    textobject = c.beginText(50, 750)
    textobject.setFont("Helvetica", 12)
    for line in content.splitlines():
        textobject.textLine(line)
    c.drawText(textobject)
    c.save()

    return content, filepath

# === Gradio UI with updated background color and styling ===
with gr.Blocks(css="""
body {
    background: #2A7B9B;
    background: linear-gradient(90deg,rgba(42, 123, 155, 1) 0%, rgba(87, 199, 133, 1) 50%, rgba(237, 221, 83, 1) 100%);
    font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
    color: #f0f0f0;
    margin: 0;
    min-height: 100vh;
}
.gr-box {
    background-color: #ffffff;
    border-radius: 12px;
    box-shadow: 0 4px 10px rgba(0,0,0,0.1);
    padding: 1rem 1.5rem;
    margin-bottom: 1rem;
}
.gr-button {
    border-radius: 10px;
    font-weight: 600;
    background-color: #0d6efd;
    color: white;
    transition: background-color 0.3s ease;
}
.gr-button:hover {
    background-color: #084298 !important;
    color: #fff !important;
}
h1, h2, h3, label, .gr-tab-label {
    color: #212529 !important;
    font-weight: 700;
}
/* Tabs title */
.gr-tabs .gr-tab-label {
    color: #0d6efd !important;
    font-weight: 700 !important;
    background-color: #e7f1ff !important;
    border-radius: 8px 8px 0 0 !important;
    padding: 10px 16px !important;
    margin-right: 4px !important;
    user-select: none;
    cursor: pointer;
}
.gr-tabs .gr-tab-label[aria-selected="true"] {
    background-color: #0d6efd !important;
    color: white !important;
    box-shadow: 0 4px 6px rgba(13, 110, 253, 0.4);
}
/* Input elements */
.gr-textbox, .gr-dropdown, .gr-file {
    border: 1.5px solid #ced4da;
    border-radius: 8px;
    padding: 0.5rem;
    font-size: 1rem;
}
""") as demo:

    gr.Markdown("""
    <div style="text-align: center; margin-bottom: 20px;">
        <h1>🩺 Doctor TWIN – Your Virtual Healthcare Companion</h1>
        <p>A lightweight AI assistant for quick medical Q&A and OTC prescription generation</p>
    </div>
    """)

    with gr.Tabs():
        with gr.TabItem("💬 Doctor Twin Advice"):
            with gr.Row():
                with gr.Column(scale=1, min_width=300):
                    gr.Markdown("### 👤 Enter Your Query", elem_classes=["gr-box"])
                    user_input = gr.Textbox(
                        lines=4,
                        label="Describe your symptom or question",
                        placeholder="e.g., I have a mild fever and fatigue",
                        elem_classes=["gr-textbox"]
                    )
                    category = gr.Dropdown(
                        label="Symptom Category",
                        choices=["General", "Skin", "Mental Health", "Respiratory", "Digestive"],
                        value="General",
                        elem_classes=["gr-dropdown"]
                    )
                    submit = gr.Button("Get Advice", elem_classes=["gr-button"])

                with gr.Column(scale=2):
                    gr.Markdown("### 🤖 Doctor Twin Says", elem_classes=["gr-box"])
                    output = gr.Textbox(label="AI Response", lines=8, interactive=False, show_copy_button=True, elem_classes=["gr-textbox"])

            submit.click(fn=doctor_twin_light, inputs=[user_input, category], outputs=output, show_progress=True)

        with gr.TabItem("📄 OTC Prescription"):
            with gr.Row():
                with gr.Column(scale=1):
                    patient_name = gr.Textbox(label="Patient Name", placeholder="e.g., John Doe", elem_classes=["gr-textbox"])
                    symptoms_input = gr.Textbox(label="Symptoms", placeholder="e.g., cough, runny nose", elem_classes=["gr-textbox"])
                    gen_button = gr.Button("Generate Prescription", elem_classes=["gr-button"])

                with gr.Column(scale=2):
                    prescription_output = gr.Textbox(label="Generated Prescription", lines=10, interactive=False, show_copy_button=True, elem_classes=["gr-textbox"])
                    pdf_file = gr.File(label="Download Prescription PDF", elem_classes=["gr-file"])

            gen_button.click(
                fn=generate_otc_prescription,
                inputs=[patient_name, symptoms_input],
                outputs=[prescription_output, pdf_file],
                show_progress=True
            )

demo.launch()