#!/usr/bin/env python3
"""
Bacteria Detection Application
Hugging Face Spaces - Gradio Interface
"""

import gradio as gr
import cv2
import numpy as np
from ultralytics import YOLO
from PIL import Image
import os

# Model path
MODEL_PATH = "bacteria_model.pt"

# Class names
CLASS_NAMES = {
    0: 'G- Cocci',
    1: 'G+ Cocci',
    2: 'G- Bacilli',
    3: 'G+ Bacilli'
}

# Colors (BGR)
CLASS_COLORS = {
    0: (255, 150, 100),
    1: (100, 100, 255),
    2: (255, 200, 100),
    3: (100, 200, 255),
}

# Load model
print("Loading model...")
model = YOLO(MODEL_PATH)
print("Model loaded successfully!")

def process_image(image):
    """
    Process image and detect bacteria
    
    Args:
        image: PIL Image or numpy array
        
    Returns:
        tuple: (processed_image, detection_info)
    """
    # Convert PIL to OpenCV format
    if isinstance(image, Image.Image):
        image = np.array(image)
    
    # Convert RGB to BGR for OpenCV
    img_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
    
    # Predict
    results = model.predict(
        source=img_bgr,
        conf=0.25,
        iou=0.45,
        verbose=False
    )
    
    # Initialize counters
    detection_count = 0
    class_counts = {name: 0 for name in CLASS_NAMES.values()}
    
    # Draw boxes
    for r in results:
        detection_count = len(r.boxes)
        for box in r.boxes:
            # Get coordinates
            x1, y1, x2, y2 = box.xyxy[0].cpu().numpy().astype(int)
            cls_id = int(box.cls[0])
            conf = float(box.conf[0])
            
            # Update class count
            class_name = CLASS_NAMES[cls_id]
            class_counts[class_name] += 1
            
            # Draw box
            color = CLASS_COLORS.get(cls_id, (0, 255, 0))
            cv2.rectangle(img_bgr, (x1, y1), (x2, y2), color, 2)
            
            # Draw label
            label = f"{class_name} {conf:.0%}"
            (w, h), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)
            cv2.rectangle(img_bgr, (x1, y1 - h - 10), (x1 + w, y1), color, -1)
            cv2.putText(img_bgr, label, (x1, y1 - 5),
                       cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)
    
    # Convert back to RGB for display
    img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
    
    # Create detection summary
    summary = f"**Total Detections: {detection_count} bacteria**\n\n"
    if detection_count > 0:
        summary += "**Class Distribution:**\n"
        for class_name, count in class_counts.items():
            if count > 0:
                summary += f"- {class_name}: {count}\n"
    else:
        summary = "No bacteria detected."
    
    return img_rgb, summary

# Create Gradio interface
with gr.Blocks(title="🦠 Bacteria Detection", theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # 🦠 Bacteria Detection System
        
        Upload a microscope image to detect and classify bacteria.
        
        **Classes:**
        - 🔴 G- Cocci (Gram-negative cocci)
        - 🔵 G+ Cocci (Gram-positive cocci)
        - 🟠 G- Bacilli (Gram-negative bacilli)
        - 🟡 G+ Bacilli (Gram-positive bacilli)
        """
    )
    
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(
                type="pil",
                label="Upload Microscope Image",
                height=400
            )
            detect_btn = gr.Button("🔬 Detect Bacteria", variant="primary", size="lg")
            
            gr.Markdown(
                """
                ### Instructions:
                1. Upload a microscope image (JPG, PNG)
                2. Click "Detect Bacteria" button
                3. View results on the right
                
                **Note:** The model detects bacteria based on Gram staining and shape.
                """
            )
        
        with gr.Column():
            output_image = gr.Image(
                type="numpy",
                label="Detection Results",
                height=400
            )
            output_text = gr.Markdown(label="Detection Summary")
    
    # Examples (if available)
    example_files = []
    if os.path.exists("examples"):
        example_files = [os.path.join("examples", f) for f in os.listdir("examples") if f.endswith(('.jpg', '.png'))]
    
    if example_files:
        gr.Markdown("### 📸 Try These Examples:")
        gr.Examples(
            examples=example_files,
            inputs=input_image,
        )
    
    # Event handler
    detect_btn.click(
        fn=process_image,
        inputs=input_image,
        outputs=[output_image, output_text]
    )
    
    gr.Markdown(
        """
        ---
        ### 📊 Model Information
        - **Architecture:** YOLOv8
        - **Classes:** 4 (G- Cocci, G+ Cocci, G- Bacilli, G+ Bacilli)
        - **Training Dataset:** Clinical Bacteria Detection Dataset
        - **Confidence Threshold:** 25%
        
        ### ℹ️ About
        This application uses deep learning to detect and classify bacteria in microscope images
        based on Gram staining (Gram-positive/negative) and shape (cocci/bacilli).
        """
    )

# Launch
if __name__ == "__main__":
    demo.launch()