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

# Load model
model = YOLO('models/best.pt')

# Class mapping sesuai dengan mobile app
CLASS_NAMES = {
    0: 'amblas',
    1: 'bergelombang', 
    2: 'berlubang',
    3: 'retak_buaya'
}

def detect_road_damage(image, confidence_threshold=0.5):
    """
    Deteksi kerusakan jalan dari gambar
    """
    try:
        # Convert PIL to numpy array
        if isinstance(image, Image.Image):
            image = np.array(image)
        
        # Run inference
        results = model(image, conf=confidence_threshold)
        
        detections = []
        annotated_image = image.copy()
        
        for result in results:
            boxes = result.boxes
            if boxes is not None:
                for box in boxes:
                    # Extract detection info
                    x1, y1, x2, y2 = box.xyxy[0].cpu().numpy()
                    confidence = float(box.conf[0])
                    class_id = int(box.cls[0])
                    class_name = CLASS_NAMES.get(class_id, 'unknown')
                    
                    # Calculate dimensions (estimasi)
                    width_pixels = x2 - x1
                    height_pixels = y2 - y1
                    
                    # Estimasi ukuran dalam cm (asumsi 1 pixel = 0.1 cm)
                    width_cm = width_pixels * 0.1
                    depth_cm = height_pixels * 0.1
                    
                    detection = {
                        'class': class_name,
                        'confidence': confidence,
                        'bbox': [float(x1), float(y1), float(x2), float(y2)],
                        'width_cm': width_cm,
                        'depth_cm': depth_cm
                    }
                    detections.append(detection)
                    
                    # Draw bounding box
                    cv2.rectangle(annotated_image, 
                                (int(x1), int(y1)), (int(x2), int(y2)), 
                                (0, 255, 0), 2)
                    
                    # Add label
                    label = f"{class_name}: {confidence:.2f}"
                    cv2.putText(annotated_image, label,
                              (int(x1), int(y1-10)),
                              cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
        
        # Prepare response
        response = {
            'detections': detections,
            'count': len(detections),
            'status': 'success'
        }
        
        return annotated_image, json.dumps(response, indent=2)
        
    except Exception as e:
        error_response = {
            'error': str(e),
            'status': 'error'
        }
        return image, json.dumps(error_response, indent=2)

# Gradio interface
with gr.Blocks(title="VGTec Road Damage Detector") as demo:
    gr.Markdown("# 🛣️ VGTec Road Damage Detection API")
    gr.Markdown("Upload gambar jalan untuk mendeteksi kerusakan (amblas, bergelombang, berlubang, retak_buaya)")
    
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(type="pil", label="Upload Gambar Jalan")
            confidence_slider = gr.Slider(
                minimum=0.1, 
                maximum=1.0, 
                value=0.5, 
                step=0.1,
                label="Confidence Threshold"
            )
            detect_btn = gr.Button("🔍 Deteksi Kerusakan", variant="primary")
        
        with gr.Column():
            output_image = gr.Image(label="Hasil Deteksi")
            output_json = gr.Code(label="JSON Response", language="json")
    
    # Event handler
    detect_btn.click(
        fn=detect_road_damage,
        inputs=[input_image, confidence_slider],
        outputs=[output_image, output_json]
    )
    
    # Examples
    gr.Examples(
        examples=[
            ["examples/berlubang.jpg", 0.5],
            ["examples/retak_buaya.jpg", 0.6],
        ],
        inputs=[input_image, confidence_slider],
        outputs=[output_image, output_json],
        fn=detect_road_damage,
        cache_examples=True
    )

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