import gradio as gr
from transformers import pipeline
from PIL import Image
import torch
import os
import spaces

# Initialize the model pipeline
print("Loading MedGemma model...")
pipe = pipeline(
    "image-text-to-text",
    model="google/medgemma-4b-it",
    torch_dtype=torch.bfloat16,
    device="cuda" if torch.cuda.is_available() else "cpu",
)
print("Model loaded successfully!")

@spaces.GPU()
def analyze_xray(image, custom_prompt=None):
    """
    Analyze X-ray image using MedGemma model
    """
    if image is None:
        return "Please upload an X-ray image first."
    
    try:
        # Use custom prompt if provided, otherwise use default
        if custom_prompt and custom_prompt.strip():
            prompt_text = custom_prompt.strip()
        else:
            prompt_text = "Describe this X-ray in detail, including any abnormalities or notable findings."
        
        messages = [
            {
                "role": "system",
                "content": [{"type": "text", "text": "You are an expert radiologist with years of experience in interpreting medical images."}]
            },
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": prompt_text},
                    {"type": "image", "image": image},
                ]
            }
        ]
        
        # Generate analysis
        output = pipe(text=messages, max_new_tokens=300)
        result = output[0]["generated_text"][-1]["content"]
        
        return result
        
    except Exception as e:
        return f"Error analyzing image: {str(e)}"

def load_sample_image():
    """Load the sample X-ray image if it exists"""
    sample_path = "./images/Chest_Xray_PA_3-8-2010.png"
    if os.path.exists(sample_path):
        return Image.open(sample_path)
    return None

# Create Gradio interface
with gr.Blocks(
    theme=gr.themes.Soft(),
    title="AI X-ray Analysis System",
    css="""
    .header {
        text-align: center;
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
        color: white;
        padding: 2rem;
        border-radius: 10px;
        margin-bottom: 2rem;
    }
    .warning {
        background-color: #fff3cd;
        border: 1px solid #ffeaa7;
        border-radius: 8px;
        padding: 1rem;
        margin: 1rem 0;
        color: #856404;
    }
    .gradio-container {
        max-width: 1200px;
        margin: auto;
    }
    """
) as demo:
    
    # Header
    gr.HTML("""
        <div class="header">
            <h1>🩻 AI X-ray Analysis System</h1>
            <p>Advanced medical image analysis powered by Google's MedGemma AI</p>
        </div>
    """)
    
    # Warning disclaimer
    gr.HTML("""
        <div class="warning">
            <strong>⚠️ Medical Disclaimer:</strong> This AI tool is for educational and research purposes only. 
            It should not be used as a substitute for professional medical diagnosis or treatment. 
            Always consult qualified healthcare professionals for medical advice.
        </div>
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 📤 Upload X-ray Image")
            
            # Image input
            image_input = gr.Image(
                label="X-ray Image",
                type="pil",
                height=400,
                sources=["upload", "clipboard"]
            )
            
            # Sample image button
            sample_btn = gr.Button(
                "📋 Load Sample Image", 
                variant="secondary",
                size="sm"
            )
            
            # Custom prompt input
            gr.Markdown("### 💬 Custom Analysis Prompt (Optional)")
            custom_prompt = gr.Textbox(
                label="Custom Prompt",
                placeholder="Enter specific questions about the X-ray (e.g., 'Focus on the heart area' or 'Look for signs of pneumonia')",
                value = "Describe this X-ray",
                lines=3,
                max_lines=5
            )
            
            # Analyze button
            analyze_btn = gr.Button(
                "🔍 Analyze X-ray", 
                variant="primary",
                size="lg"
            )
        
        with gr.Column(scale=1):
            gr.Markdown("### 📊 Analysis Results")
            
            # Output text
            output_text = gr.Textbox(
                label="AI Analysis Report",
                lines=28,
                max_lines=100,
                show_copy_button=True,
                placeholder="Upload an X-ray image and click 'Analyze X-ray' to see the AI analysis results here..."
            )
            
            # Quick action buttons
            with gr.Row():
                clear_btn = gr.Button("🗑️ Clear", variant="secondary", size="sm")
                copy_btn = gr.Button("📋 Copy Results", variant="secondary", size="sm")
    
    # Example prompts section
    gr.Markdown("### 💡 Example Prompts")
    with gr.Row():
        example_prompts = [
            "Describe this X-ray in detail, including any abnormalities or notable findings.",
            "Focus on the lung fields and identify any signs of infection or disease.",
            "Examine the heart size and shape. Is the cardiac silhouette normal?",
            "Look for any signs of fractures or bone abnormalities.",
            "Analyze the overall image quality and positioning."
        ]
        
        for i, prompt in enumerate(example_prompts):
            gr.Button(
                f"Example {i+1}",
                size="sm"
            ).click(
                lambda p=prompt: p,
                outputs=custom_prompt
            )
    
    # Event handlers
    def clear_all():
        return None, "", ""
    
    sample_btn.click(
        fn=load_sample_image,
        outputs=image_input
    )
    
    analyze_btn.click(
        fn=analyze_xray,
        inputs=[image_input, custom_prompt],
        outputs=output_text
    )
    
    clear_btn.click(
        fn=clear_all,
        outputs=[image_input, custom_prompt, output_text]
    )
    
    # Auto-analyze when image is uploaded (optional)
    image_input.change(
        fn=lambda img: analyze_xray(img) if img is not None else "",
        inputs=image_input,
        outputs=output_text
    )

# Launch the app
if __name__ == "__main__":
    print("Starting Gradio interface...")
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,  # Set to True if you want to create a public link
        show_error=True,
        favicon_path=None
    )