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

# Load the YOLOv8 model
model = YOLO('yolov8n.pt')

def process_image(image):
    """
    Process a single image for object detection
    """
    results = model(image)
    # Get detection information
    boxes = results[0].boxes
    detection_info = []
    for box in boxes:
        class_id = int(box.cls[0])
        class_name = results[0].names[class_id]
        confidence = float(box.conf[0])
        detection_info.append(f"{class_name}: {confidence:.2%}")
    
    return Image.fromarray(results[0].plot()), "\n".join(detection_info)

def process_video(video_path):
    """
    Process video for object detection
    """
    with tempfile.NamedTemporaryFile(delete=False, suffix='.mp4') as temp_file:
        output_path = temp_file.name

    cap = cv2.VideoCapture(video_path)
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = int(cap.get(cv2.CAP_PROP_FPS))
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
    
    detection_summary = []
    frame_count = 0
    
    try:
        while cap.isOpened():
            ret, frame = cap.read()
            if not ret:
                break
                
            frame_count += 1
            results = model(frame)
            
            # Collect detection information for this frame
            if frame_count % int(fps) == 0:  # Sample every second
                for box in results[0].boxes:
                    class_id = int(box.cls[0])
                    class_name = results[0].names[class_id]
                    detection_summary.append(class_name)
            
            annotated_frame = results[0].plot()
            out.write(annotated_frame)
            
    finally:
        cap.release()
        out.release()
    
    # Create summary of detected objects
    if detection_summary:
        from collections import Counter
        counts = Counter(detection_summary)
        summary = "\n".join([f"{obj}: {count} occurrences" for obj, count in counts.most_common()])
    else:
        summary = "No objects detected"
    
    return output_path, summary

def detect_objects(media):
    """
    Unified function to handle both image and video inputs
    """
    if media is None:
        return None, None, None, "Please upload an image or video to begin detection.", gr.update(visible=True), gr.update(visible=False)
    
    try:
        if isinstance(media, str) and media.lower().endswith(('.mp4', '.avi', '.mov')):
            output_video, detection_summary = process_video(media)
            return (None, output_video, detection_summary,
                    "✅ Video processing complete! Check the detection summary below.", 
                    gr.update(visible=False), gr.update(visible=True))
        else:
            if isinstance(media, str):
                image = cv2.imread(media)
                image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            else:
                image = media
            processed_image, detection_info = process_image(image)
            return (processed_image, None, detection_info,
                    "✅ Image processing complete! Check the detections below.", 
                    gr.update(visible=True), gr.update(visible=False))
    except Exception as e:
        return None, None, None, f"❌ Error: {str(e)}", gr.update(visible=False), gr.update(visible=False)

# Custom CSS for styling
custom_css = """
#app-container {
    max-width: 1200px;
    margin: 0 auto;
    padding: 20px;
}

#logo-img {
    max-height: 100px;
    margin-bottom: 20px;
}

.upload-box {
    border: 2px dashed #ccc;
    padding: 20px;
    text-align: center;
    border-radius: 8px;
    background-color: #f8f9fa;
    margin: 20px 0;
}

.results-container {
    background-color: #ffffff;
    border-radius: 8px;
    padding: 15px;
    box-shadow: 0 2px 4px rgba(0,0,0,0.1);
    margin-top: 20px;
}

.detection-info {
    background-color: #f8f9fa;
    padding: 15px;
    border-radius: 8px;
    margin-top: 10px;
    font-family: monospace;
}

.center {
    display: flex;
    justify-content: center;
    align-items: center;
    margin-bottom: 1rem;
}
"""

# Create Gradio interface
with gr.Blocks(css=custom_css) as demo:
    with gr.Column(elem_id="app-container"):
        # Logo and Header
        with gr.Column(elem_classes="center"):
            gr.Image("logo-h.png", 
                    show_label=False, 
                    container=False,
                    elem_id="logo-img",
                    height=100)
        
        gr.Markdown("# 🔍 Object Detection")
        
        # Upload Section
        with gr.Column(elem_classes="upload-box"):
            gr.Markdown("### 📤 Upload your file")
            input_media = gr.File(
                label="Drag and drop or click to upload (Images: jpg, jpeg, png | Videos: mp4, avi, mov)",
                file_types=["image", "video"]
            )
        
        # Status Message
        status_text = gr.Textbox(
            label="Status",
            value="Waiting for upload...",
            interactive=False
        )
        
        # Detection Information
        detection_info = gr.Textbox(
            label="Detection Results",
            elem_classes="detection-info",
            interactive=False
        )
        
        # Results Section
        with gr.Column(elem_classes="results-container"):
            with gr.Row():
                with gr.Column(visible=False) as image_column:
                    output_image = gr.Image(label="Detected Objects")
                with gr.Column(visible=False) as video_column:
                    output_video = gr.Video(label="Processed Video")
    
        # Handle file upload
        input_media.upload(
            fn=detect_objects,
            inputs=[input_media],
            outputs=[
                output_image,
                output_video,
                detection_info,
                status_text,
                image_column,
                video_column
            ]
        )

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