"""Gradio demo for ObjectModel-v1 — a compact, NMS-free COCO object detector."""

from __future__ import annotations

import os
import time

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces  # noqa: E402  — must precede torch / any CUDA-touching import

import numpy as np  # noqa: E402
import torch  # noqa: E402
import gradio as gr  # noqa: E402
from PIL import Image, ImageDraw, ImageFont  # noqa: E402
from huggingface_hub import hf_hub_download  # noqa: E402

from objectmodel_v1.boxes import box_cxcywh_to_xyxy  # noqa: E402
from objectmodel_v1.config import load_config  # noqa: E402
from objectmodel_v1.model import build_model  # noqa: E402

MODEL_ID = "bench-labs/objectmodel-v1"
CHECKPOINT_FILE = "objectmodel_v1_best.pt"
CONFIG_FILE = "objectmodel_v1.yaml"

# The training dataloader maps COCO category ids -> contiguous labels by ascending
# category id, so label i is the i-th COCO category in id order.
COCO_CLASSES = [
    "person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck",
    "boat", "traffic light", "fire hydrant", "stop sign", "parking meter", "bench",
    "bird", "cat", "dog", "horse", "sheep", "cow", "elephant", "bear", "zebra",
    "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee",
    "skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove",
    "skateboard", "surfboard", "tennis racket", "bottle", "wine glass", "cup",
    "fork", "knife", "spoon", "bowl", "banana", "apple", "sandwich", "orange",
    "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch",
    "potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse",
    "remote", "keyboard", "cell phone", "microwave", "oven", "toaster", "sink",
    "refrigerator", "book", "clock", "vase", "scissors", "teddy bear", "hair drier",
    "toothbrush",
]

# ImageNet statistics, matching the training-time preprocessing.
MEAN = torch.tensor([0.485, 0.456, 0.406], dtype=torch.float32)[:, None, None]
STD = torch.tensor([0.229, 0.224, 0.225], dtype=torch.float32)[:, None, None]

# ---------------------------------------------------------------------------
# Model — loaded once at module scope and moved to CUDA eagerly (ZeroGPU rule 2).
# ---------------------------------------------------------------------------
_config = load_config(CONFIG_FILE)
_checkpoint_path = hf_hub_download(MODEL_ID, CHECKPOINT_FILE)

model = build_model(_config)
_state = torch.load(_checkpoint_path, map_location="cpu", weights_only=False)
model.load_state_dict(_state.get("ema", _state.get("model", _state)))
del _state
model.eval().to("cuda")

INPUT_SIZE = int(model.spec.input_size)
NUM_PARAMS = sum(p.numel() for p in model.parameters())
print(f"ObjectModel-v1 loaded: {NUM_PARAMS/1e6:.1f}M params, input {INPUT_SIZE}px", flush=True)


def letterbox(image: Image.Image, size: int) -> tuple[torch.Tensor, float, int, int]:
    """Resize keeping aspect ratio and pad to a square canvas (eval-time transform)."""
    width, height = image.size
    ratio = min(size / width, size / height)
    resized_width = max(1, round(width * ratio))
    resized_height = max(1, round(height * ratio))
    resized = image.resize((resized_width, resized_height), Image.Resampling.BILINEAR)
    offset_x = (size - resized_width) // 2
    offset_y = (size - resized_height) // 2
    canvas = Image.new("RGB", (size, size), (114, 114, 114))
    canvas.paste(resized, (offset_x, offset_y))
    array = np.asarray(canvas, dtype=np.float32).copy() / 255.0
    tensor = torch.from_numpy(array).permute(2, 0, 1)
    tensor = (tensor - MEAN) / STD
    return tensor, ratio, offset_x, offset_y


def palette(index: int) -> tuple[int, int, int]:
    """Stable, well-spread colour per class index."""
    hue = (index * 0.6180339887) % 1.0
    i = int(hue * 6)
    f = hue * 6 - i
    q, t = 1 - f, f
    table = [(1, t, 0), (q, 1, 0), (0, 1, t), (0, q, 1), (t, 0, 1), (1, 0, q)]
    r, g, b = table[i % 6]
    return (int(70 + 185 * r), int(70 + 185 * g), int(70 + 185 * b))


def _font(size: int):
    for name in ("DejaVuSans-Bold.ttf", "DejaVuSans.ttf"):
        try:
            return ImageFont.truetype(name, size)
        except Exception:
            continue
    return ImageFont.load_default()


def draw_detections(image: Image.Image, detections: list[dict]) -> Image.Image:
    canvas = image.convert("RGB").copy()
    draw = ImageDraw.Draw(canvas)
    scale = max(canvas.width, canvas.height) / 900.0
    thickness = max(2, round(3 * scale))
    font = _font(max(12, round(15 * scale)))
    pad = max(2, round(3 * scale))
    # Draw large boxes first so small-object labels end up on top of them.
    ordered = sorted(
        detections,
        key=lambda d: (d["box"][2] - d["box"][0]) * (d["box"][3] - d["box"][1]),
        reverse=True,
    )
    for det in ordered:
        x0, y0, x1, y1 = det["box"]
        colour = palette(det["label_index"])
        draw.rectangle([x0, y0, x1, y1], outline=colour, width=thickness)
        caption = f"{det['label']} {det['score']:.2f}"
        left, top, right, bottom = draw.textbbox((0, 0), caption, font=font)
        text_w, text_h = right - left, bottom - top
        box_w, box_h = text_w + 2 * pad, text_h + 2 * pad
        text_x = min(x0, max(0.0, canvas.width - box_w))
        text_y = y0 - box_h if y0 - box_h >= 0 else y0
        draw.rectangle([text_x, text_y, text_x + box_w, text_y + box_h], fill=colour)
        draw.text((text_x + pad - left, text_y + pad - top), caption, fill=(20, 20, 20), font=font)
    return canvas


@spaces.GPU(duration=15)
def detect(
    image: Image.Image,
    confidence: float = 0.35,
    max_detections: int = 100,
) -> tuple[Image.Image, list[list], str]:
    """Detect COCO objects in an image with ObjectModel-v1.

    Args:
        image: input photograph to run detection on.
        confidence: minimum score (0-1) a detection must reach to be kept.
        max_detections: hard cap on how many boxes are returned.

    Returns:
        The image with boxes drawn, a table of detections, and a short summary.
    """
    if image is None:
        raise gr.Error("Please upload an image first.")

    image = image.convert("RGB")
    original_width, original_height = image.size
    tensor, ratio, offset_x, offset_y = letterbox(image, INPUT_SIZE)

    started = time.perf_counter()
    with torch.inference_mode():
        outputs = model(tensor[None].to("cuda"))
        probabilities = outputs["pred_logits"].sigmoid()[0]
        boxes = box_cxcywh_to_xyxy(outputs["pred_boxes"][0]).clamp(0.0, 1.0) * INPUT_SIZE
        scores, labels = probabilities.max(dim=-1)
        count = min(int(max_detections), scores.numel())
        scores, indices = scores.topk(count)
        keep = scores >= float(confidence)
        scores, indices = scores[keep], indices[keep]
        labels = labels[indices]
        boxes = boxes[indices]
        # Undo the letterbox transform back into original image coordinates.
        boxes[:, [0, 2]] = (boxes[:, [0, 2]] - offset_x) / ratio
        boxes[:, [1, 3]] = (boxes[:, [1, 3]] - offset_y) / ratio
        boxes[:, [0, 2]] = boxes[:, [0, 2]].clamp(0, original_width)
        boxes[:, [1, 3]] = boxes[:, [1, 3]].clamp(0, original_height)
        scores = scores.float().cpu().tolist()
        labels = labels.cpu().tolist()
        boxes = boxes.float().cpu().tolist()
    elapsed = time.perf_counter() - started

    detections = [
        {
            "label": COCO_CLASSES[label] if label < len(COCO_CLASSES) else str(label),
            "label_index": int(label),
            "score": float(score),
            "box": [float(v) for v in box],
        }
        for score, label, box in zip(scores, labels, boxes)
    ]

    annotated = draw_detections(image, detections)
    table = [
        [
            det["label"],
            round(det["score"], 3),
            round(det["box"][0]),
            round(det["box"][1]),
            round(det["box"][2]),
            round(det["box"][3]),
        ]
        for det in detections
    ]
    if not table:
        table = [["(nothing above threshold)", 0.0, 0, 0, 0, 0]]
    summary = (
        f"**{len(detections)} object(s)** above {float(confidence):.2f} confidence · "
        f"forward pass {elapsed * 1000:.0f} ms · NMS-free (300 queries)"
    )
    return annotated, table, summary


CSS = """
#col-container { max-width: 1200px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks(title="ObjectModel-v1 Detection") as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            """
            # ObjectModel-v1 — compact NMS-free object detection

            A 40.8M-parameter clean-room detector from Bench Labs: compressed global latent
            memory + box-scaled local pyramid sampling, trained from scratch on COCO
            (AP 0.358). No anchors, no NMS — a fixed set of 300 queries.

            [Model card](https://huggingface.co/bench-labs/objectmodel-v1)
            """
        )
        with gr.Row():
            with gr.Column(scale=1):
                image_input = gr.Image(label="Input image", type="pil", height=380)
                confidence = gr.Slider(
                    label="Confidence threshold",
                    minimum=0.05,
                    maximum=0.95,
                    step=0.01,
                    value=0.35,
                )
                run_button = gr.Button("Detect objects", variant="primary")
                with gr.Accordion("Advanced settings", open=False):
                    max_detections = gr.Slider(
                        label="Max detections",
                        minimum=1,
                        maximum=300,
                        step=1,
                        value=100,
                        info="Top-k queries kept before the confidence filter.",
                    )
            with gr.Column(scale=1):
                image_output = gr.Image(label="Detections", type="pil", height=380)
                summary_output = gr.Markdown()
                detections_output = gr.Dataframe(
                    headers=["label", "score", "x0", "y0", "x1", "y1"],
                    label="Detections",
                    wrap=True,
                    row_count=(1, "dynamic"),
                )

        gr.Examples(
            examples=[
                ["examples/skateboarder_rail.jpg"],
                ["examples/motorcycle_street.jpg"],
                ["examples/cafe_interior.jpg"],
                ["examples/living_room_blue_couch.jpg"],
                ["examples/girl_with_dog.jpg"],
                ["examples/pizza_board.jpg"],
            ],
            inputs=[image_input],
            outputs=[image_output, detections_output, summary_output],
            fn=detect,
            cache_examples=True,
            cache_mode="lazy",
        )

    run_button.click(
        fn=detect,
        inputs=[image_input, confidence, max_detections],
        outputs=[image_output, detections_output, summary_output],
        api_name="detect",
    )

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)