"""DETR 车辆检测模块：加载模型并输出车辆类别、置信度和检测框。"""

from __future__ import annotations

from typing import Any

import torch
from PIL import Image
from transformers import DetrForObjectDetection, DetrImageProcessor


class DetrVehicleDetector:
    def __init__(self, model_name: str, confidence: float, vehicle_labels: set[str]):
        self.confidence = confidence
        self.vehicle_labels = vehicle_labels
        self.device = self._select_device()
        self.processor = DetrImageProcessor.from_pretrained(model_name)
        self.model = DetrForObjectDetection.from_pretrained(model_name)
        self.model.to(self.device)
        self.model.eval()

    @staticmethod
    def _select_device() -> torch.device:
        if torch.backends.mps.is_available():
            return torch.device("mps")
        return torch.device("cpu")

    @property
    def device_name(self) -> str:
        return self.device.type

    @torch.no_grad()
    def detect(self, frame_rgb: Any) -> list[dict[str, Any]]:
        image = Image.fromarray(frame_rgb)
        inputs = self.processor(images=image, return_tensors="pt")
        inputs = {key: value.to(self.device) for key, value in inputs.items()}

        outputs = self.model(**inputs)
        return self._detections_from_outputs(image, outputs)

    @torch.no_grad()
    def inspect_tokens(self, frame_rgb: Any) -> dict[str, Any]:
        image = Image.fromarray(frame_rgb)
        inputs = self.processor(images=image, return_tensors="pt")
        inputs = {key: value.to(self.device) for key, value in inputs.items()}

        features, object_queries_list = self.model.model.backbone(inputs["pixel_values"], inputs["pixel_mask"])
        feature_map, mask = features[-1]
        projected_feature_map = self.model.model.input_projection(feature_map)
        tokens = projected_feature_map.flatten(2).permute(0, 2, 1)
        object_queries = object_queries_list[-1].flatten(2).permute(0, 2, 1)

        outputs = self.model(**inputs, output_hidden_states=True)
        target_sizes = torch.tensor([image.size[::-1]], device=self.device)
        results = self.processor.post_process_object_detection(
            outputs,
            target_sizes=target_sizes,
            threshold=self.confidence,
        )[0]

        token_rows = int(projected_feature_map.shape[2])
        token_cols = int(projected_feature_map.shape[3])
        sample_count = min(48, int(tokens.shape[1]))
        sample_tokens = tokens[0, :sample_count, :8].detach().cpu()
        token_sequence = []
        for index, vector in enumerate(sample_tokens):
            token_sequence.append(
                {
                    "index": index,
                    "row": index // token_cols,
                    "col": index % token_cols,
                    "values": [round(float(value), 4) for value in vector.tolist()],
                    "magnitude": round(float(vector.norm()), 4),
                }
            )

        detections = []
        for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
            label_name = self.model.config.id2label[label.item()]
            if label_name not in self.vehicle_labels:
                continue

            x1, y1, x2, y2 = box.detach().cpu().numpy().astype(int).tolist()
            detections.append(
                {
                    "label": label_name,
                    "score": round(float(score.detach().cpu()), 4),
                    "box": [x1, y1, x2, y2],
                }
            )

        encoder_last_hidden_state = getattr(outputs, "encoder_last_hidden_state", None)
        last_hidden_state = getattr(outputs, "last_hidden_state", None)
        return {
            "image_size": {"width": image.size[0], "height": image.size[1]},
            "pixel_values_shape": list(inputs["pixel_values"].shape),
            "pixel_mask_shape": list(inputs["pixel_mask"].shape),
            "feature_map_shape": list(feature_map.shape),
            "projected_feature_map_shape": list(projected_feature_map.shape),
            "visual_tokens_shape": list(tokens.shape),
            "position_encoding_shape": list(object_queries.shape),
            "encoder_last_hidden_state_shape": list(encoder_last_hidden_state.shape) if encoder_last_hidden_state is not None else [],
            "decoder_last_hidden_state_shape": list(last_hidden_state.shape) if last_hidden_state is not None else [],
            "logits_shape": list(outputs.logits.shape),
            "pred_boxes_shape": list(outputs.pred_boxes.shape),
            "token_grid": {"rows": token_rows, "cols": token_cols, "total": int(tokens.shape[1]), "shown": sample_count},
            "token_sequence": token_sequence,
            "detections": detections,
        }

    def _detections_from_outputs(self, image: Image.Image, outputs: Any) -> list[dict[str, Any]]:
        target_sizes = torch.tensor([image.size[::-1]], device=self.device)
        results = self.processor.post_process_object_detection(
            outputs,
            target_sizes=target_sizes,
            threshold=self.confidence,
        )[0]

        detections = []
        for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
            label_name = self.model.config.id2label[label.item()]
            if label_name not in self.vehicle_labels:
                continue

            x1, y1, x2, y2 = box.detach().cpu().numpy().astype(int).tolist()
            detections.append(
                {
                    "label": label_name,
                    "score": round(float(score.detach().cpu()), 4),
                    "box": [x1, y1, x2, y2],
                }
            )

        return detections