import gradio as gr
import cv2
import numpy as np
import torch
from ultralytics import YOLO
import time
from typing import Tuple, Dict, List, Optional, Union

# Initialize default confidence threshold
DEFAULT_CONF_THRESHOLD = 0.25

# Load YOLO model (using YOLOv8 as it's optimized for CPU)
model = YOLO('yolov8n.pt')  # using the nano version for faster CPU inference

# Set to CPU explicitly
model.to('cpu')

def detect_objects(image, conf_threshold: float = DEFAULT_CONF_THRESHOLD):
    if image is None:
        return None, "No image provided"
    
    # Convert image from RGB (Gradio) to BGR (OpenCV)
    image_cv = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
    
    # Start timer for FPS calculation
    start_time = time.time()
    
    # Run inference with user-defined confidence threshold
    results = model(image_cv, conf=conf_threshold)
    
    # Calculate FPS
    fps = 1.0 / (time.time() - start_time)
    
    # Process results
    if len(results) == 0:
        return image, "No detections found"
    
    result = results[0]
    output_image = image.copy()
    
    # Get detections
    boxes = result.boxes.cpu().numpy()
    
    detection_count = {}
    
    # Draw bounding boxes and labels
    for box in boxes:
        x1, y1, x2, y2 = map(int, box.xyxy[0])
        conf = float(box.conf[0])
        cls_id = int(box.cls[0])
        cls_name = result.names[cls_id]
        
        # Update detection count
        detection_count[cls_name] = detection_count.get(cls_name, 0) + 1
        
        # Generate a color based on class ID (for consistent colors per class)
        color = (int(hash(cls_name) % 256), int(hash(cls_name + "salt") % 256), int(hash(cls_name + "pepper") % 256))
        
        # Draw bounding box
        cv2.rectangle(output_image, (x1, y1), (x2, y2), color, 2)
        
        # Draw label background
        text = f"{cls_name}: {conf:.2f}"
        text_size, _ = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
        cv2.rectangle(output_image, (x1, y1 - text_size[1] - 5), (x1 + text_size[0], y1), color, -1)
        
        # Draw text
        cv2.putText(output_image, text, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
    
    # Create detection summary
    summary = f"FPS: {fps:.1f}\n"
    if detection_count:
        summary += "Detected objects:\n"
        for cls_name, count in sorted(detection_count.items()):
            summary += f"- {cls_name}: {count}\n"
    else:
        summary += "No objects detected"
    
    return output_image, summary

def process_webcam(image, conf_threshold: float = DEFAULT_CONF_THRESHOLD):
    """Process webcam frames for real-time detection"""
    return detect_objects(image, conf_threshold)

def process_uploaded_image(image, conf_threshold: float = DEFAULT_CONF_THRESHOLD):
    """Process uploaded image for detection"""
    return detect_objects(image, conf_threshold)

# Create custom HTML header for the app
custom_css = """
<style>
.header-container {
    background: linear-gradient(90deg, #2C3E50, #4CA1AF);
    color: white;
    padding: 1.5rem;
    border-radius: 8px;
    margin-bottom: 1.5rem;
    box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
    text-align: center;
}
.header-title {
    font-size: 2.5rem;
    font-weight: bold;
    margin-bottom: 0.5rem;
}
.header-subtitle {
    font-size: 1.2rem;
    opacity: 0.9;
}
</style>
"""

custom_header = f"""
{custom_css}
<div class="header-container">
    <div class="header-title">MyStuff Object Detection Demo</div>
    <div class="header-subtitle">Powered by YOLOv8 and Gradio</div>
</div>
"""

# Create Gradio interface - using Gradio 3.x compatible syntax
with gr.Blocks() as demo:
    gr.HTML(custom_header)
    
    # Global confidence threshold slider
    with gr.Row():
        conf_threshold = gr.Slider(
            minimum=0.1, 
            maximum=0.9, 
            value=DEFAULT_CONF_THRESHOLD, 
            step=0.05, 
            label="Confidence Threshold", 
            info="Adjust to control detection sensitivity (lower = more detections)"
        )
    
    with gr.Tabs():
        with gr.TabItem("Webcam Detection"):
            with gr.Row():
                with gr.Column():
                    webcam_input = gr.Image(label="Webcam Input")
                    webcam_button = gr.Button("Enable Webcam", variant="primary")
                with gr.Column():
                    webcam_output = gr.Image(label="Detection Output")
                    webcam_summary = gr.Textbox(label="Detection Summary")
            
            # Process button for webcam
            process_webcam_button = gr.Button("Process Webcam Image")
            process_webcam_button.click(
                process_webcam,
                inputs=[webcam_input, conf_threshold],
                outputs=[webcam_output, webcam_summary]
            )
        
        with gr.TabItem("Image Upload Detection"):
            with gr.Row():
                with gr.Column():
                    img_input = gr.Image(label="Upload Image")
                    img_button = gr.Button("Detect Objects", variant="primary")
                with gr.Column():
                    img_output = gr.Image(label="Detection Output")
                    img_summary = gr.Textbox(label="Detection Summary")
            
            img_button.click(
                process_uploaded_image,
                inputs=[img_input, conf_threshold],
                outputs=[img_output, img_summary]
            )
    
    with gr.Accordion("Usage Instructions & Notes", open=True):
        gr.Markdown("""
        ### Usage Instructions
        - **Adjust Confidence Threshold**: Use the slider to control detection sensitivity
          - **Higher values** (e.g., 0.7-0.9): Only very confident detections, fewer objects detected
          - **Lower values** (e.g., 0.1-0.3): More objects detected but may include false positives
        - **Webcam Tab:** 
          1. Click "Enable Webcam" to activate your camera
          2. Click "Process Webcam Image" to detect objects in the current frame
        - **Image Upload Tab:** Upload an image and click "Detect Objects"
        
        ### Technical Notes
        - This demo uses YOLOv8 nano model optimized for CPU usage
        - Performance will vary based on your CPU speed
        - Detection runs locally in your browser
        - The FPS counter shows real-time performance metrics
        """)

# Launch the app
if __name__ == "__main__":
    demo.launch(share=True)