app.py

"""
MedVision AI — Medical Imaging Assistant
Gradio app for Build Small Hackathon (Hugging Face Spaces)
"""

import warnings
from typing import Dict, List, Optional, Tuple

import gradio as gr
from PIL import Image

from utils import (
    LangChainMedicalChatbot,
    MedGemmaRunnable,
)

warnings.filterwarnings("ignore")

# ---------------------------------------------------------------------------
# Module-level client state
# ---------------------------------------------------------------------------
medgemma_llm = None
chatbot: Optional[LangChainMedicalChatbot] = None
model_load_error: Optional[str] = None

DISCLAIMER_HTML = """
<div class="disclaimer-banner">
  <strong>⚠️ Medical Disclaimer:</strong>
  This tool provides <em>clinical decision support only</em> — not a diagnosis.
  Always verify findings with clinical examination, correlate with patient history,
  consult specialists when needed, and follow institutional protocols.
</div>
"""

CUSTOM_CSS = """
:root {
  --bg-primary: #0f172a;
  --bg-card: #1e293b;
  --accent: #14b8a6;
  --accent-hover: #0d9488;
  --text-primary: #f1f5f9;
  --text-muted: #94a3b8;
  --user-bubble: #1d4ed8;
}

.gradio-container {
  background: var(--bg-primary) !important;
  color: var(--text-primary) !important;
  max-width: 1400px !important;
}

.app-header {
  background: linear-gradient(135deg, #0f172a 0%, #1e3a5f 50%, #0f766e 100%);
  border: 1px solid var(--accent);
  border-radius: 12px;
  padding: 1.5rem 2rem;
  margin-bottom: 1rem;
  text-align: center;
}

.app-header h1 {
  color: #f8fafc;
  font-size: 1.75rem;
  margin: 0 0 0.25rem 0;
  font-weight: 700;
}

.app-header p {
  color: var(--text-muted);
  margin: 0;
  font-size: 0.95rem;
}

.disclaimer-banner {
  background: linear-gradient(90deg, #78350f, #92400e);
  border: 1px solid #f59e0b;
  border-radius: 8px;
  padding: 0.75rem 1rem;
  color: #fef3c7;
  font-size: 0.875rem;
  margin-bottom: 1rem;
  line-height: 1.5;
}

.block, .panel {
  background: var(--bg-card) !important;
  border: 1px solid #334155 !important;
  border-radius: 10px !important;
}

label, .label-wrap span {
  color: var(--text-primary) !important;
}

button.primary {
  background: var(--accent) !important;
  border-color: var(--accent) !important;
  color: #0f172a !important;
  font-weight: 600 !important;
}

button.primary:hover {
  background: var(--accent-hover) !important;
}

button.secondary {
  border-color: #475569 !important;
  color: var(--text-muted) !important;
}

.model-status {
  padding: 0.5rem 1rem;
  border-radius: 6px;
  font-size: 0.85rem;
  margin-bottom: 0.75rem;
}

.model-status.loading {
  background: #1e3a5f;
  border: 1px solid #3b82f6;
  color: #93c5fd;
}

.model-status.ready {
  background: #064e3b;
  border: 1px solid var(--accent);
  color: #5eead4;
}

.model-status.error {
  background: #450a0a;
  border: 1px solid #ef4444;
  color: #fca5a5;
}

/* Chatbot bubbles */
.message.user {
  background: var(--user-bubble) !important;
  border-color: #2563eb !important;
}

.message.bot, .message.assistant {
  background: var(--bg-card) !important;
  border: 1px solid #334155 !important;
}

.chatbot {
  border: 1px solid var(--accent) !important;
  border-radius: 10px !important;
}

/* Findings clinical report cards */
.findings-card {
  display: flex;
  flex-wrap: wrap;
  gap: 0.5rem;
  padding: 0.5rem 0;
}

.finding-tag {
  background: #0f766e;
  border: 1px solid var(--accent);
  border-radius: 6px;
  padding: 0.35rem 0.75rem;
  font-size: 0.8rem;
  color: #ccfbf1;
  font-family: 'Courier New', monospace;
}

.memory-display textarea {
  font-family: 'Courier New', Consolas, monospace !important;
  font-size: 0.8rem !important;
  background: #0f172a !important;
  color: #a5f3fc !important;
  border: 1px solid #334155 !important;
}

footer {
  display: none !important;
}
"""


def load_model() -> str:
    """Initialize MedGemma Modal endpoint and LangChain chatbot. Returns status HTML."""
    global medgemma_llm, chatbot, model_load_error

    if chatbot is not None:
        return _status_html("ready", "✅ MedGemma client ready.")

    print("📡 Connecting to MedGemma Modal endpoint...")

    try:
        medgemma_llm = MedGemmaRunnable()
        chatbot = LangChainMedicalChatbot(medgemma_llm)
        print("✅ MedGemma client initialized")
    except Exception as exc:
        model_load_error = str(exc)
        return _status_html("error", f"❌ Failed to initialize: {exc}")

    return _status_html("ready", "✅ Connected to MedGemma API · Modal endpoint")


def _status_html(state: str, message: str) -> str:
    return f'<div class="model-status {state}">{message}</div>'


def _findings_to_label(findings: List[str]) -> Dict[str, float]:
    return {finding: 1.0 for finding in findings}


def _collect_all_findings(result: Dict, chatbot_instance: LangChainMedicalChatbot) -> Dict[str, float]:
    detected = result.get("findings", []) or []
    stats_findings = chatbot_instance.get_stats().get("Unique Findings", [])
    all_findings = list(set(detected + stats_findings))
    return _findings_to_label(all_findings)


def _format_findings_html(findings_dict: Dict[str, float]) -> str:
    if not findings_dict:
        return '<p style="color:#94a3b8;">No findings detected yet.</p>'
    tags = "".join(
        f'<span class="finding-tag">{label} ({score:.0%})</span>'
        for label, score in findings_dict.items()
    )
    return f'<div class="findings-card">{tags}</div>'


def run_analysis(
    image: Optional[Image.Image],
    image_type: str,
    chain_mode: str,
    question: str,
    history: List[Dict[str, str]],
) -> Tuple[List[Dict[str, str]], Dict[str, float], Dict, str, str, List[Dict[str, str]]]:
    """Route query to the appropriate LangChain chain and update UI state."""
    history = history or []

    if chatbot is None:
        gr.Warning("Model is still loading or failed to load. Please wait or refresh the page.")
        error_msg = model_load_error or "Model not loaded."
        history.append({"role": "user", "content": question or "(no question)"})
        history.append({"role": "assistant", "content": f"⚠️ {error_msg}"})
        return history, {}, {}, "", "", history

    if not question or not question.strip():
        gr.Warning("Please enter a clinical question before analyzing.")
        return history, _findings_to_label([]), chatbot.get_stats(), chatbot.get_memory(), "", history

    question = question.strip()

    try:
        if chain_mode == "Differential Diagnosis":
            result = chatbot.get_differential_diagnosis(question)

        elif chain_mode == "Follow-up / Text Query":
            result = chatbot.ask_followup(question)

        elif chain_mode == "Image Analysis":
            if image is None:
                gr.Warning("Image Analysis requires an uploaded medical image.")
                history.append({"role": "user", "content": question})
                history.append(
                    {"role": "assistant", "content": "⚠️ Please upload a medical image for Image Analysis mode."}
                )
                return (
                    history,
                    _findings_to_label([]),
                    chatbot.get_stats(),
                    chatbot.get_memory(),
                    _format_findings_html({}),
                    history,
                )
            result = chatbot.analyze_image(
                image=image,
                image_type=image_type,
                clinical_question=question,
            )

        elif chain_mode == "Auto (Smart Routing)" and image is not None:
            result = chatbot.analyze_image(
                image=image,
                image_type=image_type,
                clinical_question=question,
            )

        else:
            if image is not None:
                result = chatbot.chat(question, image=image)
            else:
                result = chatbot.ask_followup(question)

    except Exception as exc:
        gr.Warning(f"Inference error: {exc}")
        history.append({"role": "user", "content": question})
        history.append({"role": "assistant", "content": f"⚠️ An error occurred: {exc}"})
        return history, _findings_to_label([]), chatbot.get_stats(), chatbot.get_memory(), "", history

    response = result["response"]
    chain_used = result.get("chain_used", "unknown")
    timestamp = result.get("timestamp", "")
    assistant_msg = f"{response}\n\n---\n*Chain: {chain_used} · {timestamp}*"

    history.append({"role": "user", "content": question})
    history.append({"role": "assistant", "content": assistant_msg})

    findings_dict = _collect_all_findings(result, chatbot)
    stats = chatbot.get_stats()
    memory = chatbot.get_memory()
    findings_html = _format_findings_html(findings_dict)

    return history, findings_dict, stats, memory, findings_html, history


def reset_session() -> Tuple[List, Dict, Dict, str, str, List]:
    """Reset chatbot memory, stats, and all UI outputs."""
    if chatbot is not None:
        chatbot.reset()
    return (
        [],
        {},
        {},
        "",
        '<p style="color:#94a3b8;">Session reset. Upload an image and ask a clinical question to begin.</p>',
        [],
    )


# ---------------------------------------------------------------------------
# Build Gradio UI
# ---------------------------------------------------------------------------
model_status_html = load_model()

with gr.Blocks(
    title="MedVision AI",
    theme=gr.themes.Base(
        primary_hue="teal",
        secondary_hue="blue",
        neutral_hue="slate",
    ).set(
        body_background_fill="#0f172a",
        block_background_fill="#1e293b",
        block_border_color="#334155",
        body_text_color="#f1f5f9",
    ),
    css=CUSTOM_CSS,
) as demo:
    gr.HTML(
        """
        <div class="app-header">
          <h1>🏥 MedVision AI — Medical Imaging Assistant</h1>
          <p>Powered by MedGemma Modal API · LangChain</p>
        </div>
        """
    )
    gr.HTML(DISCLAIMER_HTML)
    gr.HTML(model_status_html)

    chat_state = gr.State([])

    with gr.Row():
        # LEFT COLUMN (40%)
        with gr.Column(scale=4):
            image_input = gr.Image(
                type="pil",
                label="Upload Medical Image",
                sources=["upload"],
            )
            image_type = gr.Dropdown(
                choices=[
                    "Chest X-ray",
                    "Brain MRI",
                    "CT Scan - Abdomen",
                    "CT Scan - Head",
                    "MRI - Spine",
                    "Ultrasound",
                    "Other Medical Image",
                ],
                value="Chest X-ray",
                label="Image Type",
            )
            chain_mode = gr.Dropdown(
                choices=[
                    "Auto (Smart Routing)",
                    "Image Analysis",
                    "Follow-up / Text Query",
                    "Differential Diagnosis",
                ],
                value="Auto (Smart Routing)",
                label="Analysis Mode",
            )
            clinical_question = gr.Textbox(
                lines=3,
                label="Clinical Question",
                placeholder="e.g. Identify abnormalities and provide differential diagnoses...",
            )
            with gr.Row():
                analyze_btn = gr.Button("🔬 Analyze", variant="primary")
                reset_btn = gr.Button("🔄 Reset Session", variant="secondary")

            with gr.Accordion("📊 Session Stats", open=False):
                stats_output = gr.JSON(label="Session Statistics")

        # RIGHT COLUMN (60%)
        with gr.Column(scale=6):
            chatbot_ui = gr.Chatbot(
                height=500,
                label="Clinical Consultation",
            )
            with gr.Accordion("🔍 Detected Findings", open=False):
                findings_label = gr.Label(label="Keyword Findings", num_top_classes=20)
                findings_html = gr.HTML(
                    value='<p style="color:#94a3b8;">Findings will appear after analysis.</p>'
                )
            with gr.Accordion("🧠 Conversation Memory", open=False):
                memory_output = gr.Textbox(
                    label="LangChain Memory (raw)",
                    lines=8,
                    interactive=False,
                    elem_classes=["memory-display"],
                )

    analyze_btn.click(
        fn=run_analysis,
        inputs=[image_input, image_type, chain_mode, clinical_question, chat_state],
        outputs=[chatbot_ui, findings_label, stats_output, memory_output, findings_html, chat_state],
    )

    reset_btn.click(
        fn=reset_session,
        outputs=[chatbot_ui, findings_label, stats_output, memory_output, findings_html, chat_state],
    )

    gr.Markdown(
        """
        ---
        **MedVision AI** · Powered by MedGemma 1.5 via Modal ·
        Endpoint configured in `utils.py`.
        """
    )

if __name__ == "__main__":
    demo.launch(share=False, server_name="0.0.0.0")