"""Result types returned by `Model.predict`.

Stdlib `@dataclass(slots=True, frozen=True)` — zero runtime overhead, no
pydantic. Every Result carries `boxes`, `labels`, `scores` in the base;
task-specific subclasses add their own fields (`masks` for segmentation,
future: `rotated_boxes` for OBB, `text` for OCR).

Tensors stay on whatever device the model was on. Call `.cpu()` before
post-processing if needed.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import TYPE_CHECKING

import torch

if TYPE_CHECKING:
    import supervision as sv


@dataclass(frozen=True, slots=True)
class Result:
    """Base result shared across all tasks.

    Attributes:
        boxes: shape `[N, 4]`, xyxy in image pixel coordinates, float32.
        labels: shape `[N]`, int64 class ids.
        scores: shape `[N]`, float32 confidence scores in [0, 1].
    """

    boxes: torch.Tensor
    labels: torch.Tensor
    scores: torch.Tensor


@dataclass(frozen=True, slots=True)
class DetectionResult(Result):
    """Object-detection result. No task-specific fields beyond the base."""

    def to_supervision(self) -> sv.Detections:
        """Convert to `supervision.Detections` for post-processing/annotation."""
        import supervision as sv

        return sv.Detections(
            xyxy=self.boxes.detach().cpu().numpy(),
            class_id=self.labels.detach().cpu().numpy(),
            confidence=self.scores.detach().cpu().numpy(),
        )


@dataclass(frozen=True, slots=True)
class SegmentationResult(Result):
    """Instance-segmentation result. Adds per-instance binary masks.

    Attributes:
        masks: shape `[N, H, W]`, boolean tensor of per-instance masks at
            the original input resolution.
    """

    masks: torch.Tensor

    def to_supervision(self) -> sv.Detections:
        """Convert to `supervision.Detections`, including masks."""
        import supervision as sv

        return sv.Detections(
            xyxy=self.boxes.detach().cpu().numpy(),
            class_id=self.labels.detach().cpu().numpy(),
            confidence=self.scores.detach().cpu().numpy(),
            mask=self.masks.detach().cpu().numpy().astype(bool),
        )
