app.py

"""
app.py — Medical Image Explainability Interface with Interactive AI Chatbot.

Features:
- Chest X-ray upload and classification (Pneumonia / Tuberculosis / Normal)
- Grad-CAM heatmap visualization with anatomical localization
- Interactive AI chatbot powered by Gemini for medical image Q&A
- Premium dark medical theme with glassmorphism and animations
"""

import json
import gradio as gr
import numpy as np
from PIL import Image
import google.generativeai as genai

from model import load_model, CLASS_NAMES
from gradcam import run_inference
from explainer import generate_explanation, generate_metadata_display

# ── Configuration ──────────────────────────────────────────────────────────────
GEMINI_API_KEY = "AIzaSyDC9Fl5Okuzamt-xZZpJmVwOlIQXib26FY"
genai.configure(api_key=GEMINI_API_KEY)
# Use available gemini-2.5-flash model
gemini_model = genai.GenerativeModel("gemini-2.5-flash")

# Load classification model at startup
print("Loading medical classification model...")
model = load_model()
print("Model loaded successfully.")

# Global state for current analysis
current_analysis = {"result": None, "explanation": None}

SYSTEM_PROMPT = """You are a Medical Image Explainability Agent. Your role is to provide 
evidence-based explanations for chest X-ray model classifications (Pneumonia, Tuberculosis, Normal) 
using Grad-CAM heatmap data and anatomical localization.

GROUNDING PROTOCOL:
1. Every claim must originate from the anatomical_focus and heatmap_data provided. 
   Do not reference lung fields or lobes not flagged by the activation map.
2. Follow the Observation → Tool Evidence → Explanation sequence.
3. Do NOT infer patient symptoms (fever, cough), clinical history, or demographics.
4. Explain Grad-CAM as back-propagating gradients to the final convolutional layer 
   to highlight pixels influencing the decision.

GRAD-CAM METHODOLOGY:
Grad-CAM computes neuron importance weights α_k^c for class c and feature map A^k:
   α_k^c = (1/Z) Σ_i Σ_j ∂Y^c / ∂A^k_ij
The heatmap is: L_GradCAM^c = ReLU(Σ_k α_k^c · A^k)

When the user asks questions, be helpful, precise, and ground all answers in the 
available image analysis data. If no image has been analyzed yet, let the user know 
they need to upload and analyze an X-ray first.

Keep responses clear, professional, and formatted with markdown."""


# ── Custom CSS ─────────────────────────────────────────────────────────────────
CUSTOM_CSS = """
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700;800&family=JetBrains+Mono:wght@400;500&display=swap');

:root {
    --bg-primary: #060a14;
    --bg-secondary: #0c1222;
    --bg-card: rgba(12, 18, 34, 0.7);
    --bg-card-hover: rgba(18, 26, 48, 0.8);
    --border-subtle: rgba(99, 102, 241, 0.1);
    --border-glow: rgba(99, 102, 241, 0.25);
    --accent-blue: #6366f1;
    --accent-purple: #8b5cf6;
    --accent-cyan: #22d3ee;
    --accent-emerald: #34d399;
    --accent-rose: #fb7185;
    --text-primary: #e2e8f0;
    --text-secondary: #94a3b8;
    --text-muted: #64748b;
}

* { font-family: 'Inter', -apple-system, BlinkMacSystemFont, sans-serif !important; }

.gradio-container {
    background: var(--bg-primary) !important;
    background-image: 
        radial-gradient(ellipse at 20% 50%, rgba(99, 102, 241, 0.06) 0%, transparent 50%),
        radial-gradient(ellipse at 80% 20%, rgba(139, 92, 246, 0.04) 0%, transparent 50%),
        radial-gradient(ellipse at 50% 80%, rgba(34, 211, 238, 0.03) 0%, transparent 50%) !important;
    min-height: 100vh;
}

/* ── Header ────────────────────────────────────────────────────────────────── */
.hero-header {
    text-align: center;
    padding: 2.5rem 2rem;
    margin-bottom: 2rem;
    background: linear-gradient(145deg, rgba(99, 102, 241, 0.06), rgba(139, 92, 246, 0.04), rgba(34, 211, 238, 0.03));
    border: 1px solid var(--border-subtle);
    border-radius: 20px;
    backdrop-filter: blur(24px);
    position: relative;
    overflow: hidden;
}

.hero-header::before {
    content: '';
    position: absolute;
    top: 0; left: 0; right: 0;
    height: 2px;
    background: linear-gradient(90deg, transparent, var(--accent-blue), var(--accent-purple), var(--accent-cyan), transparent);
    animation: shimmer 3s ease-in-out infinite;
}

@keyframes shimmer {
    0%, 100% { opacity: 0.5; }
    50% { opacity: 1; }
}

.hero-header h1 {
    font-size: 2.2rem;
    font-weight: 800;
    letter-spacing: -0.03em;
    margin: 0 0 0.4rem 0;
    background: linear-gradient(135deg, #818cf8, #a78bfa, #c084fc, #22d3ee);
    background-size: 200% 200%;
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    animation: gradient-flow 4s ease infinite;
}

@keyframes gradient-flow {
    0% { background-position: 0% 50%; }
    50% { background-position: 100% 50%; }
    100% { background-position: 0% 50%; }
}

.hero-header .subtitle {
    color: var(--text-secondary);
    font-size: 1rem;
    font-weight: 300;
    margin: 0;
    letter-spacing: 0.01em;
}

.hero-header .badge-row {
    margin-top: 1rem;
    display: flex;
    justify-content: center;
    gap: 0.75rem;
    flex-wrap: wrap;
}

.hero-header .badge {
    display: inline-flex;
    align-items: center;
    gap: 0.35rem;
    padding: 0.3rem 0.75rem;
    border-radius: 999px;
    font-size: 0.75rem;
    font-weight: 500;
    background: rgba(99, 102, 241, 0.08);
    border: 1px solid rgba(99, 102, 241, 0.15);
    color: #a5b4fc;
}

/* ── Section Labels ────────────────────────────────────────────────────────── */
.section-title {
    display: flex;
    align-items: center;
    gap: 0.5rem;
    color: var(--text-secondary) !important;
    font-weight: 600 !important;
    font-size: 0.8rem !important;
    text-transform: uppercase !important;
    letter-spacing: 0.1em !important;
    margin: 0 0 0.75rem 0 !important;
    padding-bottom: 0.5rem;
    border-bottom: 1px solid var(--border-subtle);
}

.section-title .icon {
    width: 20px;
    height: 20px;
    border-radius: 6px;
    display: flex;
    align-items: center;
    justify-content: center;
    font-size: 0.7rem;
}

/* ── Buttons ───────────────────────────────────────────────────────────────── */
.gr-button-primary {
    background: linear-gradient(135deg, #6366f1, #8b5cf6, #7c3aed) !important;
    border: none !important;
    border-radius: 12px !important;
    font-weight: 600 !important;
    padding: 0.85rem 2rem !important;
    font-size: 0.95rem !important;
    letter-spacing: 0.01em !important;
    transition: all 0.3s cubic-bezier(0.4, 0, 0.2, 1) !important;
    box-shadow: 0 4px 20px rgba(99, 102, 241, 0.25), inset 0 1px 0 rgba(255,255,255,0.1) !important;
    text-shadow: 0 1px 2px rgba(0,0,0,0.2) !important;
}

.gr-button-primary:hover {
    transform: translateY(-2px) !important;
    box-shadow: 0 8px 30px rgba(99, 102, 241, 0.4), inset 0 1px 0 rgba(255,255,255,0.15) !important;
}

.gr-button-primary:active {
    transform: translateY(0) !important;
}

/* ── Cards & Panels ────────────────────────────────────────────────────────── */
.panel, .gr-panel {
    background: var(--bg-card) !important;
    border: 1px solid var(--border-subtle) !important;
    border-radius: 14px !important;
    backdrop-filter: blur(16px) !important;
    transition: border-color 0.3s ease !important;
}

.panel:hover {
    border-color: var(--border-glow) !important;
}

/* ── Chatbot Styles ────────────────────────────────────────────────────────── */
.chatbot-container {
    border: 1px solid rgba(99, 102, 241, 0.12) !important;
    border-radius: 14px !important;
    background: rgba(8, 12, 24, 0.6) !important;
    backdrop-filter: blur(12px) !important;
}

/* ── Metadata Code Block ───────────────────────────────────────────────────── */
.metadata-panel {
    background: rgba(6, 10, 20, 0.9) !important;
    border: 1px solid rgba(34, 211, 238, 0.12) !important;
    border-radius: 12px !important;
    font-family: 'JetBrains Mono', monospace !important;
    font-size: 0.82rem !important;
}

/* ── Explanation ───────────────────────────────────────────────────────────── */
.explanation-panel {
    background: rgba(12, 18, 34, 0.5) !important;
    border: 1px solid rgba(139, 92, 246, 0.1) !important;
    border-radius: 14px !important;
    padding: 1.25rem !important;
    line-height: 1.75 !important;
    color: var(--text-primary) !important;
}

/* ── Disclaimer ────────────────────────────────────────────────────────────── */
.disclaimer {
    background: rgba(239, 68, 68, 0.04);
    border: 1px solid rgba(239, 68, 68, 0.1);
    border-radius: 12px;
    padding: 1rem 1.25rem;
    margin-top: 1.5rem;
}

.disclaimer p {
    color: #fca5a5;
    font-size: 0.78rem;
    margin: 0;
    line-height: 1.6;
}

/* ── Footer ────────────────────────────────────────────────────────────────── */
footer { display: none !important; }

.custom-footer {
    text-align: center;
    padding: 1.5rem 0 0.5rem;
    color: var(--text-muted);
    font-size: 0.72rem;
    letter-spacing: 0.02em;
}

/* ── Responsive ────────────────────────────────────────────────────────────── */
@media (max-width: 768px) {
    .hero-header h1 { font-size: 1.5rem; }
    .hero-header { padding: 1.5rem 1rem; }
}
"""


# ── Core Functions ─────────────────────────────────────────────────────────────

def classify_and_explain(image):
    """Main pipeline: classify image → generate heatmap → build explanation."""
    if image is None:
        current_analysis["result"] = None
        current_analysis["explanation"] = None
        return None, {}, "", ""

    if isinstance(image, np.ndarray):
        image = Image.fromarray(image)

    result, overlay = run_inference(model, image)

    label_dict = result["all_probabilities"]
    metadata = generate_metadata_display(result)
    explanation = generate_explanation(result)

    # Store for chatbot context
    current_analysis["result"] = result
    current_analysis["explanation"] = explanation

    return overlay, label_dict, metadata, explanation


def chatbot_respond(message, history):
    """Interactive chatbot powered by Gemini with analysis context."""
    if not message.strip():
        return ""

    # Build context from current analysis
    context_block = ""
    if current_analysis["result"]:
        r = current_analysis["result"]
        context_block = (
            "\nCURRENT ANALYSIS DATA (use this to ground your responses):\n"
            f"- Prediction: {r['prediction']}\n"
            f"- Confidence: {r['confidence'] * 100:.1f}%\n"
            f"- Anatomical Focus: {r['anatomical_focus']}\n"
            f"- Heatmap Intensity: {r['heatmap_intensity']}\n"
            f"- Methodology: {r['methodology']}\n"
            f"- All Probabilities: {json.dumps(r['all_probabilities'])}\n"
        )
    else:
        context_block = "\nNo image has been analyzed yet. Ask the user to upload and analyze an X-ray first.\n"

    # Build full prompt with system instructions + context + conversation history
    full_prompt = SYSTEM_PROMPT + "\n" + context_block + "\nCONVERSATION HISTORY:\n"
    for msg in history:
        role = msg.get("role", "")
        content = msg.get("content", "")
        if role == "user":
            full_prompt += f"User: {content}\n"
        elif role == "assistant":
            full_prompt += f"Assistant: {content}\n"

    full_prompt += f"User: {message}\nAssistant:"

    try:
        response = gemini_model.generate_content(full_prompt)
        return response.text
    except Exception as e:
        return f"⚠️ Error communicating with Gemini API: {str(e)}"


# ── Build Gradio Interface ─────────────────────────────────────────────────────

with gr.Blocks(css=CUSTOM_CSS, title="Medical Image Explainability Agent") as demo:

    # Hero Header
    gr.HTML("""
    <div class="hero-header">
        <h1>🏥 Medical Image Explainability Agent</h1>
        <p class="subtitle">AI-powered chest X-ray analysis with Grad-CAM visual explanations & interactive chatbot</p>
        <div class="badge-row">
            <span class="badge">🔬 Grad-CAM Heatmaps</span>
            <span class="badge">🧠 Gemini AI Chatbot</span>
            <span class="badge">🫁 Anatomical Localization</span>
            <span class="badge">📊 3-Class Classification</span>
        </div>
    </div>
    """)

    # ── Row 1: Upload + Heatmap + Classification ──
    gr.HTML('<p class="section-title"><span class="icon">📷</span> Image Analysis</p>')

    with gr.Row(equal_height=True):
        with gr.Column(scale=1):
            img_input = gr.Image(
                type="pil",
                label="Upload Chest X-Ray",
                elem_id="xray-input",
                height=340,
            )
            analyze_btn = gr.Button(
                "🔬 Analyze X-Ray",
                variant="primary",
                elem_id="analyze-btn",
            )

        with gr.Column(scale=1):
            heatmap_output = gr.Image(
                label="Grad-CAM Activation Overlay",
                elem_id="heatmap-output",
                height=340,
                interactive=False,
            )
            label_output = gr.Label(
                label="Prediction Probabilities",
                elem_id="classification-label",
            )

    # ── Row 2: Metadata + Explanation side by side ──
    with gr.Row(equal_height=True):
        with gr.Column(scale=1):
            gr.HTML('<p class="section-title"><span class="icon">📋</span> Agent Metadata</p>')
            metadata_output = gr.Code(
                label="Classification Metadata (JSON)",
                language="json",
                elem_id="metadata-panel",
                elem_classes=["metadata-panel"],
                interactive=False,
            )
        with gr.Column(scale=2):
            gr.HTML('<p class="section-title"><span class="icon">🧠</span> Structured Explanation</p>')
            explanation_output = gr.Markdown(
                elem_id="explanation-panel",
                elem_classes=["explanation-panel"],
            )

    # ── Row 3: Interactive Chatbot ──
    gr.HTML('<p class="section-title"><span class="icon">💬</span> AI Medical Assistant — Ask questions about the analysis</p>')

    chatbot = gr.Chatbot(
        label="Medical Image Q&A",
        elem_id="chatbot-main",
        elem_classes=["chatbot-container"],
        height=380,
        placeholder="Upload and analyze an X-ray, then ask me questions about the findings...",
    )

    with gr.Row():
        chat_input = gr.Textbox(
            placeholder="Ask about the classification, heatmap regions, or methodology...",
            label="Your Message",
            elem_id="chat-input",
            scale=5,
        )
        chat_send = gr.Button("Send ➤", variant="primary", scale=1, elem_id="chat-send-btn")

    clear_btn = gr.Button("🗑️ Clear Chat", variant="secondary", elem_id="clear-chat-btn")

    # Disclaimer
    gr.HTML("""
    <div class="disclaimer">
        <p>⚠️ <strong>Disclaimer:</strong> This tool is for educational and research purposes only.
        It is <em>not</em> a substitute for professional medical diagnosis. All analysis is strictly limited to
        visual evidence within the provided X-ray image and the model's learned feature representations.
        No clinical history, symptoms, or demographic data are inferred or implied.</p>
    </div>
    """)

    gr.HTML('<p class="custom-footer">Medical Image Explainability Agent · Powered by Grad-CAM & Gemini AI</p>')

    # ── Event Wiring ──

    analyze_btn.click(
        fn=classify_and_explain,
        inputs=[img_input],
        outputs=[heatmap_output, label_output, metadata_output, explanation_output],
    )

    img_input.change(
        fn=classify_and_explain,
        inputs=[img_input],
        outputs=[heatmap_output, label_output, metadata_output, explanation_output],
    )

    def handle_chat(message, history):
        if not message.strip():
            return history, ""
        history = history or []
        
        # Bulletproof normalizer to handle whatever format Gradio passes
        normalized_history = []
        for msg in history:
            if isinstance(msg, dict):
                normalized_history.append({"role": msg.get("role", "user"), "content": msg.get("content", "")})
            elif hasattr(msg, "role"):
                normalized_history.append({"role": getattr(msg, "role", "user"), "content": getattr(msg, "content", "")})
            elif isinstance(msg, (list, tuple)) and len(msg) == 2:
                if msg[0]: normalized_history.append({"role": "user", "content": str(msg[0])})
                if msg[1]: normalized_history.append({"role": "assistant", "content": str(msg[1])})

        bot_response = chatbot_respond(message, normalized_history)
        
        normalized_history.append({"role": "user", "content": message})
        normalized_history.append({"role": "assistant", "content": bot_response})
        
        return normalized_history, ""

    chat_send.click(
        fn=handle_chat,
        inputs=[chat_input, chatbot],
        outputs=[chatbot, chat_input],
    )

    chat_input.submit(
        fn=handle_chat,
        inputs=[chat_input, chatbot],
        outputs=[chatbot, chat_input],
    )

    clear_btn.click(fn=lambda: [], outputs=[chatbot])


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