import spaces  # must be imported before torch

import torch
import gradio as gr
from PIL import Image, ImageDraw, ImageFont
from transformers import AutoImageProcessor, AutoModelForObjectDetection

MODEL_ID = "joshlyman/rtdetrv2-floorplancad-doors"

processor = AutoImageProcessor.from_pretrained(MODEL_ID)
model = AutoModelForObjectDetection.from_pretrained(MODEL_ID)
model.eval()

COLORS = {
    "single_door": "#e6194b",
    "double_door": "#3cb44b",
    "sliding_door": "#4363d8",
}


@spaces.GPU(duration=1)
def _noop():  # ZeroGPU requires >=1 decorated function; never called (model runs on CPU)
    pass


def detect_doors(image: Image.Image, confidence_threshold: float = 0.5):
    """Detect single, double, and sliding doors in a floor-plan image and draw the boxes."""
    if image is None:
        return None, "Upload a floor-plan image first."

    image = image.convert("RGB")
    inputs = processor(images=image, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)

    target_sizes = torch.tensor([image.size[::-1]])
    results = processor.post_process_object_detection(
        outputs, threshold=confidence_threshold, target_sizes=target_sizes
    )[0]

    annotated = image.copy()
    draw = ImageDraw.Draw(annotated)
    try:
        font = ImageFont.truetype("DejaVuSans-Bold.ttf", 16)
    except OSError:
        font = ImageFont.load_default()

    rows = []
    for score, label_id, box in zip(results["scores"], results["labels"], results["boxes"]):
        label = model.config.id2label[label_id.item()]
        conf = score.item()
        x0, y0, x1, y1 = [round(v) for v in box.tolist()]
        color = COLORS.get(label, "#f58231")

        draw.rectangle([x0, y0, x1, y1], outline=color, width=3)
        caption = f"{label} {conf:.2f}"
        text_bbox = draw.textbbox((x0, y0), caption, font=font)
        draw.rectangle(
            [x0, text_bbox[1] - 2, text_bbox[2] + 4, text_bbox[3] + 2], fill=color
        )
        draw.text((x0 + 2, y0), caption, fill="white", font=font)

        rows.append([label, f"{conf:.2%}", f"({x0}, {y0}, {x1}, {y1})"])

    if not rows:
        summary = "No doors detected above this confidence threshold."
    else:
        summary = f"{len(rows)} door(s) detected."

    return annotated, gr.Dataframe(
        headers=["type", "confidence", "box (x0, y0, x1, y1)"], value=rows
    ), summary


with gr.Blocks(title="Floor Plan Door Detector") as demo:
    gr.Markdown(
        "# Floor Plan Door Detector\n"
        f"RT-DETRv2 fine-tuned on floor-plan CAD drawings — "
        f"[{MODEL_ID}](https://huggingface.co/{MODEL_ID}) — "
        "detects **single**, **double**, and **sliding** doors. Upload a floor-plan "
        "image (or an exported CAD/PDF page as an image) to see the detections."
    )
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(type="pil", label="Floor plan")
            threshold = gr.Slider(
                minimum=0.05, maximum=0.95, value=0.5, step=0.05,
                label="Confidence threshold",
            )
            run_btn = gr.Button("Detect doors", variant="primary")
        with gr.Column():
            image_output = gr.Image(type="pil", label="Detections")
            summary_output = gr.Textbox(label="Summary", interactive=False)
            table_output = gr.Dataframe(
                headers=["type", "confidence", "box (x0, y0, x1, y1)"],
                label="Detected doors",
            )

    run_btn.click(
        fn=detect_doors,
        inputs=[image_input, threshold],
        outputs=[image_output, table_output, summary_output],
        api_name="detect_doors",
    )
    image_input.change(
        fn=detect_doors,
        inputs=[image_input, threshold],
        outputs=[image_output, table_output, summary_output],
    )

demo.launch(mcp_server=True)