"""RT-DETR (Real-Time DEtection TRansformer) adapter.

Backbone: `transformers.AutoModelForObjectDetection` on
`CondadosAI/rtdetr_r50vd` (mirror of `PekingU/rtdetr_r50vd_coco_o365`).

    predict flow
    ────────────
    inputs
       │  io.load_images(strict_rgb=...)
       ▼
    [PIL RGB] x N
       │  chunk into batches of `batch_size`
       ▼
    each batch
       │  processor(images=batch, return_tensors="pt")
       │  → {pixel_values: Tensor[B, 3, H, W], ...}
       │  .to(device)
       ▼
    with torch.inference_mode():
        outputs = model(**batch_inputs)
       │
       ▼
    processor.post_process_object_detection(
        outputs, target_sizes=[(h, w), ...], threshold=0.5
    ) → list[dict{boxes, labels, scores}]
       │
       ▼
    DetectionResult per image
"""

from __future__ import annotations

from collections.abc import Sequence
from typing import Any

import torch
from transformers import AutoImageProcessor, AutoModelForObjectDetection

from acaua.io import ImageInput, load_images
from acaua.model import Model
from acaua.results import DetectionResult, Result

# Confidence below this is dropped from the returned DetectionResult.
# Matches transformers' conventional default. Users can post-filter the
# returned tensors if they want a different threshold.
_DEFAULT_THRESHOLD = 0.5


class RTDETRAdapter(Model):
    """RT-DETR object-detection adapter."""

    def __init__(
        self,
        *,
        repo_id: str,
        processor: Any,
        model: Any,
        device: torch.device,
    ) -> None:
        self.repo_id = repo_id
        self._processor = processor
        self._model = model
        self.device = device

    @classmethod
    def _load(
        cls,
        *,
        name: str,
        revision: str,
        device: torch.device,
    ) -> RTDETRAdapter:
        processor = AutoImageProcessor.from_pretrained(name, revision=revision)  # type: ignore[no-untyped-call]
        model = AutoModelForObjectDetection.from_pretrained(name, revision=revision)
        model.eval()
        model.to(device)
        return cls(repo_id=name, processor=processor, model=model, device=device)

    def predict(
        self,
        inputs: ImageInput | Sequence[ImageInput],
        *,
        batch_size: int = 8,
        strict_rgb: bool = True,
    ) -> Result | list[Result]:
        pil_images, was_sequence = load_images(inputs, strict_rgb=strict_rgb)

        results: list[Result] = []
        for start in range(0, len(pil_images), batch_size):
            batch = pil_images[start : start + batch_size]
            target_sizes = [(img.height, img.width) for img in batch]
            batch_inputs = self._processor(images=batch, return_tensors="pt")
            batch_inputs = {k: v.to(self.device) for k, v in batch_inputs.items()}
            with torch.inference_mode():
                outputs = self._model(**batch_inputs)
            per_image = self._processor.post_process_object_detection(
                outputs,
                target_sizes=target_sizes,
                threshold=_DEFAULT_THRESHOLD,
            )
            for det in per_image:
                results.append(
                    DetectionResult(
                        boxes=det["boxes"].to(torch.float32),
                        labels=det["labels"].to(torch.int64),
                        scores=det["scores"].to(torch.float32),
                    )
                )

        return results if was_sequence else results[0]
