# app.py - Main Gradio application
import gradio as gr
import torch
from transformers import AutoProcessor, AutoModelForImageTextToText
from PIL import Image
import os
import logging

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Model configuration
MODEL_ID = "google/medgemma-4b-it"

# Global variables for model and processor
model = None
processor = None

def load_model():
    """Load model and processor with error handling"""
    global model, processor
    
    try:
        logger.info(f"Loading model: {MODEL_ID}")
        
        # Check if CUDA is available
        device = "cuda" if torch.cuda.is_available() else "cpu"
        logger.info(f"Using device: {device}")
        
        # Load model with appropriate settings for Spaces
        model = AutoModelForImageTextToText.from_pretrained(
            MODEL_ID,
            torch_dtype=torch.bfloat16 if device == "cuda" else torch.float32,
            device_map="auto" if device == "cuda" else None,
            trust_remote_code=True,
            low_cpu_mem_usage=True
        )
        
        processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
        
        logger.info("Model loaded successfully!")
        return True
        
    except Exception as e:
        logger.error(f"Error loading model: {str(e)}")
        return False

# Initialize model at startup
model_loaded = load_model()

def analyze_medical_image(image, clinical_question, patient_history=""):
    """
    Analyze medical image with clinical context
    """
    global model, processor
    
    # Check if model is loaded
    if not model_loaded or model is None or processor is None:
        return "❌ Model not loaded. Please try refreshing the page or contact support."
    
    if image is None:
        return "⚠️ Please upload a medical image first."
    
    if not clinical_question.strip():
        return "⚠️ Please provide a clinical question."
    
    try:
        # Prepare the conversation with proper medical context
        messages = [
            {
                "role": "system",
                "content": [{"type": "text", "text": "You are MedGemma, an expert medical AI assistant specialized in medical image analysis. Provide detailed, structured analysis while emphasizing that this is for educational purposes only and should not replace professional medical diagnosis. Be thorough but clear in your explanations."}]
            }
        ]
        
        # Build user message content
        user_content = []
        
        # Add patient history if provided
        if patient_history.strip():
            user_content.append({"type": "text", "text": f"Patient History: {patient_history}\n\n"})
        
        # Add the clinical question
        user_content.append({"type": "text", "text": f"Clinical Question: {clinical_question}"})
        
        # Add the image
        user_content.append({"type": "image", "image": image})
        
        messages.append({
            "role": "user",
            "content": user_content
        })
        
        # Process inputs
        inputs = processor.apply_chat_template(
            messages, 
            add_generation_prompt=True, 
            tokenize=True,
            return_dict=True, 
            return_tensors="pt"
        )
        
        # Move to appropriate device
        device = next(model.parameters()).device
        dtype = next(model.parameters()).dtype
        inputs = {k: v.to(device) for k, v in inputs.items()}
        
        input_len = inputs["input_ids"].shape[-1]
        
        # Generate response with appropriate settings
        with torch.inference_mode():
            generation = model.generate(
                **inputs, 
                max_new_tokens=1500,
                do_sample=True,
                temperature=0.3,  # Lower temperature for more focused medical analysis
                top_p=0.95,
                repetition_penalty=1.1,
                pad_token_id=processor.tokenizer.eos_token_id
            )
            generation = generation[0][input_len:]
        
        # Decode response
        response = processor.decode(generation, skip_special_tokens=True)
        
        # Clean up response
        response = response.strip()
        
        # Add structured disclaimer
        disclaimer = """
        
---
### ⚠️ MEDICAL DISCLAIMER
**This analysis is for educational and research purposes only.**
- This AI assistant is not a substitute for professional medical advice
- Always consult qualified healthcare professionals for diagnosis and treatment
- Do not make medical decisions based solely on this analysis
- In case of medical emergency, contact emergency services immediately
---
        """
        
        return response + disclaimer
        
    except Exception as e:
        logger.error(f"Error in analyze_medical_image: {str(e)}")
        return f"❌ Error processing request: {str(e)}\n\nPlease try again or contact support if the issue persists."

# Create Gradio interface
def create_interface():
    with gr.Blocks(
        title="MedGemma Medical Image Analysis", 
        theme=gr.themes.Soft(),
        css="""
        .gradio-container {
            max-width: 1200px !important;
        }
        .disclaimer {
            background-color: #fef2f2;
            border: 1px solid #fecaca;
            border-radius: 8px;
            padding: 16px;
            margin: 16px 0;
        }
        """
    ) as demo:
        
        # Header
        gr.Markdown("""
        # 🏥 MedGemma Medical Image Analysis
        
        **Advanced Medical AI Assistant powered by Google's MedGemma-4B**
        
        This tool can analyze various medical imaging modalities including:
        - 🫁 **Chest X-rays** - Pneumonia, COVID-19, lung pathology
        - 🧠 **CT Scans** - Brain, chest, abdomen imaging
        - 🔬 **Histopathology** - Microscopic tissue analysis
        - 👁️ **Ophthalmology** - Retinal imaging, eye conditions
        - 🩺 **Dermatology** - Skin lesions and conditions
        """)
        
        # Warning banner
        with gr.Row():
            gr.Markdown("""
            <div class="disclaimer">
            ⚠️ <strong>IMPORTANT MEDICAL DISCLAIMER</strong><br>
            This tool is for <strong>educational and research purposes only</strong>. 
            Do not upload real patient data or use for actual medical diagnosis. 
            Always consult qualified healthcare professionals.
            </div>
            """)
        
        with gr.Row():
            # Left column - Inputs
            with gr.Column(scale=1):
                gr.Markdown("## 📤 Upload & Configure")
                
                image_input = gr.Image(
                    label="Medical Image", 
                    type="pil",
                    height=350,
                    sources=["upload", "clipboard"]
                )
                
                clinical_question = gr.Textbox(
                    label="Clinical Question *",
                    placeholder="Examples:\n• Describe the findings in this chest X-ray\n• What pathological changes are visible?\n• Provide differential diagnosis based on imaging\n• Identify any abnormalities present",
                    lines=4,
                    max_lines=6
                )
                
                patient_history = gr.Textbox(
                    label="Patient History (Optional)",
                    placeholder="Example: 65-year-old male presenting with chronic cough, shortness of breath, and chest pain. History of smoking for 30 years.",
                    lines=3,
                    max_lines=5
                )
                
                with gr.Row():
                    clear_btn = gr.Button("🗑️ Clear All", variant="secondary")
                    analyze_btn = gr.Button("🔍 Analyze Image", variant="primary", size="lg")
                
                # Model status
                gr.Markdown(f"""
                **Model Status:** {'✅ Loaded' if model_loaded else '❌ Not Loaded'}  
                **Model:** {MODEL_ID}  
                **Device:** {'CUDA' if torch.cuda.is_available() else 'CPU'}
                """)
                
            # Right column - Output
            with gr.Column(scale=1):
                gr.Markdown("## 📋 Medical Analysis Results")
                
                output = gr.Textbox(
                    label="AI Medical Analysis",
                    lines=25,
                    max_lines=35,
                    show_copy_button=True,
                    placeholder="Upload an image and ask a clinical question to get started..."
                )
        
        # Example cases section
        gr.Markdown("## 💡 Example Use Cases")
        
        with gr.Accordion("Click to see example cases", open=False):
            examples = gr.Examples(
                examples=[
                    [
                        "https://upload.wikimedia.org/wikipedia/commons/c/c8/Chest_Xray_PA_3-8-2010.png",
                        "Analyze this chest X-ray and describe any abnormal findings. Comment on the heart size, lung fields, and overall chest anatomy.",
                        "Adult patient presenting with respiratory symptoms including cough and shortness of breath."
                    ],
                    [
                        None,
                        "What pathological changes are visible in this medical image? Provide a structured analysis including anatomical observations and potential diagnoses.",
                        "Patient with acute onset symptoms"
                    ],
                    [
                        None,
                        "Perform a systematic review of this imaging study. Include: 1) Technical quality assessment, 2) Normal anatomical structures, 3) Abnormal findings, 4) Clinical significance.",
                        ""
                    ],
                    [
                        None,
                        "Compare the findings in this image to normal anatomy. What are the key differences and what might they suggest clinically?",
                        "Follow-up imaging for known condition"
                    ]
                ],
                inputs=[image_input, clinical_question, patient_history],
                label="Click any example to load it"
            )
        
        # Event handlers
        analyze_btn.click(
            fn=analyze_medical_image,
            inputs=[image_input, clinical_question, patient_history],
            outputs=output,
            show_progress=True
        )
        
        def clear_all():
            return None, "", "", ""
        
        clear_btn.click(
            fn=clear_all,
            outputs=[image_input, clinical_question, patient_history, output]
        )
        
        # Footer information
        gr.Markdown("""
        ---
        ### 📚 About MedGemma
        
        MedGemma is Google's specialized medical AI model trained on medical imaging and clinical text. 
        It excels at:
        - Multi-modal medical image analysis
        - Clinical reasoning and differential diagnosis
        - Structured medical reporting
        - Educational medical content generation
        
        **Supported Image Types:** JPEG, PNG, TIFF, DICOM (converted)  
        **Max Image Size:** 10MB  
        **Optimal Resolution:** 896x896 pixels (auto-resized)
        
        ### 🔒 Privacy & Data Policy
        - **No data storage**: Images and text are processed in real-time and not saved
        - **No patient data**: Use only synthetic, anonymized, or educational images  
        - **Educational use**: This tool is designed for learning and research purposes
        
        ### 📞 Support
        For technical issues or questions, please create an issue in the [Hugging Face Space repository](https://huggingface.co/spaces).
        
        **Model**: Google MedGemma-4B | **Framework**: Transformers + Gradio | **License**: Apache 2.0
        """)
    
    return demo

# Launch the app
if __name__ == "__main__":
    demo = create_interface()
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=True
    )