import gradio as gr
from ultralytics import YOLO
from PIL import Image
import os
import shutil

# Define model path
MODEL_PATH = "data.pth"
TEMP_MODEL_PATH = "data.pt"

# Load the model
print(f"Initializing YOLO model from {MODEL_PATH}...")
if not os.path.exists(MODEL_PATH):
    raise FileNotFoundError(f"Model file {MODEL_PATH} not found!")

# Copy data.pth to data.pt inside the running Space container to satisfy YOLO's extension requirement
if not os.path.exists(TEMP_MODEL_PATH):
    print(f"Copying {MODEL_PATH} to {TEMP_MODEL_PATH} to satisfy YOLO's extension checks...")
    shutil.copy2(MODEL_PATH, TEMP_MODEL_PATH)
    
model = YOLO(TEMP_MODEL_PATH)
print("Model loaded successfully!")


def predict(image):
    if image is None:
        return None, "<div style='color: gray; text-align: center;'>Please upload an image to start analysis.</div>"
        
    print("Analyzing image...")
    # Run YOLO prediction on the PIL image
    results = model(image)
    
    detected_objects = []
    annotated_image = image
    
    for r in results:
        # Get annotated image with bounding boxes
        annotated_image = Image.fromarray(r.plot())
        
        # Collect and print identified objects
        for c in r.boxes.cls:
            name = model.names[int(c)]
            detected_objects.append(name)
            # Print to stdout as requested (useful for logging)
            print(f"Identified: {name}")
            
    if not detected_objects:
        objects_html = "<div style='color: gray; text-align: center; font-size: 16px; font-weight: 500;'>No objects identified.</div>"
    else:
        # Create a visually rich premium badge UI for detected objects
        objects_html = "<div style='display: flex; flex-wrap: wrap; gap: 8px; justify-content: center;'>"
        for item in sorted(list(set(detected_objects))):
            count = detected_objects.count(item)
            objects_html += f'''
            <span style="
                background: linear-gradient(135deg, #14b8a6, #0d9488);
                color: white;
                padding: 6px 16px;
                border-radius: 9999px;
                font-weight: 600;
                font-size: 14px;
                box-shadow: 0 4px 6px -1px rgba(20, 184, 166, 0.2), 0 2px 4px -2px rgba(20, 184, 166, 0.2);
                border: 1px solid rgba(255, 255, 255, 0.1);
                transition: all 0.2s ease-in-out;
            ">{item} <span style="background: rgba(255, 255, 255, 0.25); padding: 1px 6px; border-radius: 9999px; font-size: 11px; margin-left: 4px;">x{count}</span></span>
            '''
        objects_html += "</div>"
        
    return annotated_image, objects_html

# Customize premium theme
custom_theme = gr.themes.Soft(
    primary_hue="teal",
    secondary_hue="cyan",
    neutral_hue="slate"
).set(
    body_background_fill="*neutral_50",
    block_background_fill="*white",
    block_border_width="1px",
    block_label_text_color="*neutral_500"
)

# Set up Gradio UI
demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="Upload Image"),
    outputs=[
        gr.Image(type="pil", label="Detected Objects (Visual Output)"),
        gr.HTML(label="Identified Objects List")
    ],
    title="🌌 Premium Image Object Identifier",
    description="Upload an image to identify objects using the custom-loaded YOLO model. Bounding boxes will be drawn on the image and identified objects will be listed.",
    theme=custom_theme,
    css=".gradio-container { max-width: 900px; margin: auto; padding: 20px; }"
)

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