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

from transformers import (
    AutoImageProcessor,
    AutoModelForObjectDetection
)

# ---------------------------------
# Device
# ---------------------------------
device = "cuda" if torch.cuda.is_available() else "cpu"

# ---------------------------------
# Load Detection Model
# ---------------------------------
MODEL_ID = "justjuu/rtdetr-v2-license-plate-detection"

processor = AutoImageProcessor.from_pretrained(
    MODEL_ID,
    use_fast=True
)

model = AutoModelForObjectDetection.from_pretrained(MODEL_ID)
model.to(device)
model.eval()

# ---------------------------------
# Load OCR Model
# ---------------------------------
reader = easyocr.Reader(
    ["en"],
    gpu=torch.cuda.is_available()
)

# ---------------------------------
# OCR Helper (MULTI-LINE SAFE)
# ---------------------------------
def extract_license_plate_text(image, box):
    x1, y1, x2, y2 = map(int, box.tolist())
    w, h = image.size

    # Smart Padding
    pad_w = (x2 - x1) * 0.05
    pad_h = (y2 - y1) * 0.05
    crop_box = (
        max(0, int(x1 - pad_w)), max(0, int(y1 - pad_h)),
        min(w, int(x2 + pad_w)), min(h, int(y2 + pad_h))
    )

    crop = np.array(image.crop(crop_box))
    # Optimization: canvas_size=512 speed up for small crops
    result = reader.readtext(crop, allowlist='ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789 ', canvas_size=512, text_threshold=0.6)

    return " ".join([res[1] for res in result]) if result else None


# ---------------------------------
# Calculate Dynamic Drawing Parameters
# ---------------------------------
def get_drawing_params(img_height, img_width):
    """
    Calculate font scale, thickness, and padding based on image dimensions.
    This ensures consistent visual appearance across different image sizes.
    """
    # Use the larger dimension as reference for scaling
    ref_size = max(img_height, img_width)
    
    # Base values calibrated for a 1200px image
    base_ref = 1200
    
    # Scale factor relative to base reference
    scale = ref_size / base_ref
    
    # Font scale: 0.6 at 1200px, scales proportionally
    # Minimum 0.15 for very small images, max 2.0 for large
    font_scale = max(0.15, min(2.0, 0.6 * scale))
    
    # Text thickness: 2 at 1200px, minimum 1
    thickness = max(1, int(round(2 * scale)))
    
    # Padding: 5 at 1200px, minimum 1
    padding = max(1, int(round(5 * scale)))
    
    # Box thickness: 3 at 1200px, minimum 1
    box_thickness = max(1, int(round(3 * scale)))
    
    return font_scale, thickness, padding, box_thickness


# ---------------------------------
# Draw Label with Background
# ---------------------------------
def draw_label(image, text, x1, y1, img_height, img_width):
    font = cv2.FONT_HERSHEY_SIMPLEX
    font_scale, thickness, padding, _ = get_drawing_params(img_height, img_width)

    text_color = (255, 255, 255)
    bg_color = (0, 128, 0)

    (text_w, text_h), baseline = cv2.getTextSize(
        text, font, font_scale, thickness
    )

    # Offset above the bounding box, scaled to image size
    offset = max(5, int(10 * (max(img_height, img_width) / 720)))
    y_text = max(y1 - offset, text_h + padding)

    x_bg1 = x1
    y_bg1 = y_text - text_h - padding
    x_bg2 = x1 + text_w + padding * 2
    y_bg2 = y_text + baseline + padding

    cv2.rectangle(
        image,
        (x_bg1, y_bg1),
        (x_bg2, y_bg2),
        bg_color,
        -1
    )

    cv2.putText(
        image,
        text,
        (x1 + padding, y_text),
        font,
        font_scale,
        text_color,
        thickness,
        cv2.LINE_AA
    )

# ---------------------------------
# Inference Function
# ---------------------------------
def detect_license_plate(image: Image.Image):
    image = image.convert("RGB")

    inputs = processor(images=image, return_tensors="pt")
    inputs = {k: v.to(device) for k, v in inputs.items()}

    with torch.no_grad():
        outputs = model(**inputs)

    target_sizes = torch.tensor([image.size[::-1]]).to(device)

    results = processor.post_process_object_detection(
        outputs,
        target_sizes=target_sizes,
        threshold=0.6
    )[0]

    image_np = np.array(image)
    annotated = image_np.copy()
    img_height, img_width = annotated.shape[:2]
    
    # Get dynamic drawing parameters based on image size
    _, _, _, box_thickness = get_drawing_params(img_height, img_width)

    recognized_texts = []

    for box in results["boxes"]:
        box = box.cpu().numpy()
        x1, y1, x2, y2 = box.astype(int)

        plate_text = extract_license_plate_text(image, box)
        recognized_texts.append(plate_text)

        # Draw bounding box with dynamic thickness
        cv2.rectangle(
            annotated,
            (x1, y1),
            (x2, y2),
            (0, 255, 0),
            box_thickness
        )

        if plate_text:
            draw_label(annotated, plate_text, x1, y1, img_height, img_width)

    return Image.fromarray(annotated), recognized_texts

# ---------------------------------
# Gradio App
# ---------------------------------
description = """
## 🔍 Overview
This application detects vehicle license plates from images and automatically extracts the text using a modern, transformer-based computer vision pipeline. The demo is built as a **production-ready Gradio app** and deployed on **Hugging Face Spaces**.
## ⚙️ How It Works
- 🧠 **RT-DETR v2** performs high-accuracy, real-time **license plate detection**
- ✂️ Detected license plate regions are **cropped with padding**
- 🔤 **EasyOCR** recognizes characters from each detected plate
- 🖼️ Bounding boxes and recognized text are **cleanly visualized** on the image with readable labels
## 📤 How to Use
1. 📸 Upload an image containing a vehicle
2. 🧠 The model automatically detects license plates
3. 🔤 License plate text is extracted using OCR
4. ✅ The annotated image and recognized text are displayed
"""
demo = gr.Interface(
    fn=detect_license_plate,
    inputs=gr.Image(type="pil", label="Upload Vehicle Image"),
    outputs=[
        gr.Image(label="Detected License Plate"),
        gr.Textbox(label="Recognized Plate Text")
    ],
    title="🚗 License Plate Recognition (RT-DETR v2)",
    description=description,
    examples=[
        ["images/example1.jpg"],
        ["images/example2.jpg"],
    ],
    cache_examples=True
)

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