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

# Load YOLOv8 Model (Nano for speed)
model = YOLO('yolov8n.pt') 

# Load EasyOCR (English)
# This will download the OCR models on first run
reader = easyocr.Reader(['en'], gpu=torch.cuda.is_available())

def detect_and_read(image, threshold):
    # 1. YOLO Detection
    results = model.predict(source=image, conf=threshold)
    
    # 2. EasyOCR Full Page Read
    # Convert PIL to numpy for EasyOCR
    img_np = np.array(image)
    ocr_results = reader.readtext(img_np)
    
    detections = []
    
    # Process YOLO boxes (UI Components)
    for r in results:
        for box in r.boxes:
            b = box.xyxy[0].cpu().numpy() 
            c = box.cls.cpu().item()
            label = model.names[int(c)]
            
            # Box coords [xmin, ymin, xmax, ymax]
            xmin, ymin, xmax, ymax = b
            
            # Link OCR text to this specific YOLO box
            box_text = ""
            for (ocr_box, text, prob) in ocr_results:
                # Calculate OCR center
                ox = (ocr_box[0][0] + ocr_box[2][0]) / 2
                oy = (ocr_box[0][1] + ocr_box[2][1]) / 2
                
                # If text is inside the UI box
                if xmin <= ox <= xmax and ymin <= oy <= ymax:
                    box_text += text + " "
            
            # Normalize coordinates for RapnssZ
            norm_box = [
                float(ymin / image.height), # ymin
                float(xmin / image.width),  # xmin
                float(ymax / image.height), # ymax
                float(xmax / image.width)   # xmax
            ]
            
            detections.append({
                "box": norm_box,
                "label": label,
                "text": box_text.strip(),
                "score": float(box.conf.cpu().item())
            })
            
    # 3. Handle Floating Text (OCR text not inside a YOLO box)
    for (ocr_box, text, prob) in ocr_results:
        already_matched = False
        for d in detections:
            if text in d["text"]:
                already_matched = True
                break
        
        if not already_matched:
            # Add as a 'text' element
            # EasyOCR box format: [[x,y], [x,y], [x,y], [x,y]]
            norm_box = [
                float(ocr_box[0][1] / image.height),
                float(ocr_box[0][0] / image.width),
                float(ocr_box[2][1] / image.height),
                float(ocr_box[2][0] / image.width)
            ]
            detections.append({
                "box": norm_box,
                "label": "text",
                "text": text,
                "score": float(prob)
            })

    return detections

# Gradio Interface
iface = gr.Interface(
    fn=detect_and_read,
    inputs=[
        gr.Image(type="pil", label="Screen Capture"),
        gr.Slider(0, 1, value=0.2, label="Radar Sensitivity (Threshold)")
    ],
    outputs=gr.JSON(label="OCR-Radar Analysis"),
    title="RapnssZ OCR-Radar // Mega v4.5",
    description="Zero-latency UI mapping using YOLOv8 and EasyOCR."
)

iface.launch()