BoundingBoxes

class BoundingBoxes:
    r"""3D bounding boxes with flexible axis conventions.

    Inspired by `torchvision.tv_tensors.BoundingBoxes`, extended to 3D.
    One instance holds $N$ boxes, each a 6-element vector whose meaning
    is determined by the
    [`format`][torchio.data.bboxes.BoundingBoxFormat].

    Args:
        data: $(N, 6)$ tensor or array.
        format: Interpretation of the 6 columns.
        labels: Optional $(N,)$ integer tensor of class labels per box.
        affine: $4 \times 4$ affine matrix. Identity if not given.
        metadata: Arbitrary metadata dict.

    Examples:
        >>> import torch, torchio as tio
        >>> boxes = tio.BoundingBoxes(
        ...     torch.tensor([[10, 20, 30, 50, 60, 70]]),
        ...     format=tio.BoundingBoxFormat.IJKIJK,
        ... )
        >>> boxes.num_boxes
        1
    """

    def __init__(
        self,
        data: Tensor | npt.ArrayLike,
        *,
        format: BoundingBoxFormat,
        labels: Tensor | None = None,
        affine: AffineMatrix | npt.ArrayLike | None = None,
        metadata: dict[str, Any] | None = None,
    ) -> None:
        self._data = self._parse_data(data)
        self._format = format
        self._labels = self._parse_labels(labels, self._data.shape[0])
        self._affine = self._parse_affine(affine)
        self._metadata: dict[str, Any] = dict(metadata) if metadata else {}

    # --- Parsing ---

    @staticmethod
    def _parse_data(data: Tensor | npt.ArrayLike) -> Tensor:
        if not isinstance(data, Tensor):
            data = torch.as_tensor(np.asarray(data), dtype=torch.float32)
        if data.ndim != 2 or data.shape[1] != 6:
            msg = f"BoundingBoxes must have shape (N, 6), got {tuple(data.shape)}"
            raise ValueError(msg)
        return data

    @staticmethod
    def _parse_labels(labels: Tensor | None, n: int) -> Tensor | None:
        if labels is None:
            return None
        if labels.shape[0] != n:
            msg = f"Expected {n} labels, got {labels.shape[0]}"
            raise ValueError(msg)
        return labels

    @staticmethod
    def _parse_affine(affine: AffineMatrix | npt.ArrayLike | None) -> AffineMatrix:
        if affine is None:
            return AffineMatrix()
        if isinstance(affine, AffineMatrix):
            return affine
        return AffineMatrix(affine)

    # --- Properties ---

    @property
    def data(self) -> Tensor:
        """$(N, 6)$ tensor of bounding box coordinates."""
        return self._data

    @property
    def format(self) -> BoundingBoxFormat:
        """Interpretation of the 6 columns."""
        return self._format

    @property
    def labels(self) -> Tensor | None:
        """$(N,)$ integer labels, or `None`."""
        return self._labels

    @property
    def affine(self) -> AffineMatrix:
        r"""$4 \times 4$ affine mapping voxel to world coordinates."""
        return self._affine

    @property
    def metadata(self) -> dict[str, Any]:
        """Arbitrary metadata dict."""
        return self._metadata

    @property
    def num_boxes(self) -> int:
        """Number of bounding boxes."""
        return self._data.shape[0]

    @property
    def device(self) -> torch.device:
        """Device the bounding box data resides on."""
        return self._data.device

    def to(self, *args: Any, **kwargs: Any) -> Self:
        """Move box coordinates and affine to a device or dtype.

        Coordinate dtype casts are applied to the box tensor. Labels
        preserve their dtype and move only to the coordinate device.
        The affine moves to supported devices but remains `float64`.

        Args:
            *args: Positional arguments forwarded to `torch.Tensor.to`.
            **kwargs: Keyword arguments forwarded to `torch.Tensor.to`.

        Returns:
            `self` (modified in-place).
        """
        self._data = self._data.to(*args, **kwargs)
        if self._labels is not None:
            self._labels = self._labels.to(device=self._data.device)
        self._affine.to(*args, **kwargs)
        return self

    # --- Methods ---

    def to_format(self, format: BoundingBoxFormat) -> Self:
        """Convert to a different bounding box format.

        Handles representation changes (corners ↔ center-size), axis
        permutations within the same type, and voxel ↔ anatomical
        conversions (using the stored affine).

        Args:
            format: Target format.

        Returns:
            New `BoundingBoxes` in the target format.
        """
        if format == self._format:
            return self._clone(format=format)

        src_axes = self._format.axes
        tgt_axes = format.axes
        src_repr = self._format.representation
        tgt_repr = format.representation

        src_type = axes_type(src_axes)
        tgt_type = axes_type(tgt_axes)

        # Step 1: normalise to corners in source axes.
        data = self._data
        if src_repr == Representation.CENTER_SIZE:
            data = _center_size_to_corners(data)

        # Step 2: axis conversion (now in corners).
        if src_axes != tgt_axes:
            if src_type == tgt_type:
                # Same family: permute + flip.
                perm, flips = get_axis_mapping(src_axes, tgt_axes)
                data = _permute_corners(data, perm, flips)
            else:
                # Cross-type: go through world coordinates.
                data = self._cross_type_corners(
                    data,
                    src_axes,
                    src_type,
                    tgt_axes,
                    tgt_type,
                )

        # Step 3: convert to target representation.
        if tgt_repr == Representation.CENTER_SIZE:
            data = _corners_to_center_size(data)

        return self._clone(data=data, format=format)

    def new_like(
        self,
        *,
        data: Tensor | npt.ArrayLike,
        labels: Tensor | None = None,
        affine: AffineMatrix | npt.ArrayLike | None = None,
    ) -> Self:
        """Create new BoundingBoxes with the same format and metadata.

        Args:
            data: New $(N, 6)$ coordinates.
            labels: New labels. If `None`, no labels.
            affine: New affine. If `None`, uses `self.affine`.
        """
        new_affine = (
            self._parse_affine(affine) if affine is not None else self._affine.clone()
        )
        return type(self)(
            data,
            format=self._format,
            labels=labels,
            affine=new_affine,
            metadata=dict(self._metadata),
        )

    # --- Internal ---

    def _clone(
        self,
        *,
        data: Tensor | None = None,
        format: BoundingBoxFormat | None = None,
    ) -> Self:
        return type(self)(
            data if data is not None else self._data.clone(),
            format=format if format is not None else self._format,
            labels=self._labels.clone() if self._labels is not None else None,
            affine=self._affine.clone(),
            metadata=dict(self._metadata),
        )

    def _cross_type_corners(
        self,
        data: Tensor,
        src_axes: str,
        src_type: AxesType,
        tgt_axes: str,
        tgt_type: AxesType,
    ) -> Tensor:
        """Convert corners between voxel and anatomical coordinate systems."""
        if src_type == AxesType.VOXEL:
            # Voxel → RAS via affine, then optionally reorder/flip within
            # anatomical.
            # First normalise voxel order to IJK.
            if src_axes != "IJK":
                perm, _ = get_axis_mapping(src_axes, "IJK")
                data = _permute_corners(data, perm, (False, False, False))
            # Apply affine to get RAS.
            data = _ijk_corners_to_world(data, self._affine)
            # The affine's orientation tells us what "world" actually is.
            world_axes = "".join(self._affine.orientation)
            if world_axes != tgt_axes:
                perm, flips = get_axis_mapping(world_axes, tgt_axes)
                data = _permute_corners(data, perm, flips)
        else:
            # Anatomical → voxel.
            # First normalise to the affine's world system.
            world_axes = "".join(self._affine.orientation)
            if src_axes != world_axes:
                perm, flips = get_axis_mapping(src_axes, world_axes)
                data = _permute_corners(data, perm, flips)
            # World → IJK via inverse affine.
            data = _world_corners_to_ijk(data, self._affine)
            # Reorder to target voxel axes if needed.
            if tgt_axes != "IJK":
                perm, _ = get_axis_mapping("IJK", tgt_axes)
                data = _permute_corners(data, perm, (False, False, False))
        return data

    # --- Dunder ---

    def __len__(self) -> int:
        return self.num_boxes

    def __repr__(self) -> str:
        return (
            f"BoundingBoxes(num_boxes={self.num_boxes},"
            f" axes={self._format.axes!r},"
            f" representation={self._format.representation.value!r})"
        )

    def __deepcopy__(self, memo: dict) -> Self:
        new = type(self)(
            self._data.clone(),
            format=self._format,
            labels=self._labels.clone() if self._labels is not None else None,
            affine=self._affine.clone(),
            metadata=dict(self._metadata),
        )
        memo[id(self)] = new
        return new

class BoundingBoxFormat:
    """Format specification for 3D bounding boxes.

    A format is defined by two components:

    - **axes**: a 3-character string specifying the coordinate system.
      Voxel axes are permutations of `"IJK"`.
      Anatomical axes use one letter from each pair
      `{R, L}`, `{A, P}`, `{S, I}` (e.g., `"RAS"`, `"LPI"`).
    - **representation**: either *corners* (two opposite corners) or
      *center_size* (center point + extent along each axis).

    Args:
        axes: 3-character axis string.
        representation: How the 6 values encode the box.

    Examples:
        >>> from torchio.data.bboxes import BoundingBoxFormat, Representation
        >>> BoundingBoxFormat("IJK", Representation.CORNERS)
        BoundingBoxFormat(axes='IJK', representation='corners')
        >>> BoundingBoxFormat("RAS", "center_size")
        BoundingBoxFormat(axes='RAS', representation='center_size')
    """

    # Predefined convenience formats, set after the class body.
    IJKIJK: BoundingBoxFormat
    IJKWHD: BoundingBoxFormat

    __slots__ = ("_axes", "_representation")

    def __init__(
        self,
        axes: str,
        representation: Representation | str = Representation.CORNERS,
    ) -> None:
        self._axes = validate_axes(axes)
        if isinstance(representation, str):
            representation = Representation(representation)
        self._representation = representation

    @property
    def axes(self) -> str:
        """3-character axis string (e.g., `'IJK'`, `'RAS'`)."""
        return self._axes

    @property
    def representation(self) -> Representation:
        """Corners or center-size."""
        return self._representation

    def __eq__(self, other: object) -> bool:
        if not isinstance(other, BoundingBoxFormat):
            return NotImplemented
        return (
            self._axes == other._axes and self._representation == other._representation
        )

    def __hash__(self) -> int:
        return hash((self._axes, self._representation))

    def __repr__(self) -> str:
        return (
            f"BoundingBoxFormat(axes={self._axes!r},"
            f" representation={self._representation.value!r})"
        )

class Representation(Enum):
    """How the six columns of a bounding box are interpreted.

    Attributes:
        CORNERS: Two corners: $(a_1, b_1, c_1, a_2, b_2, c_2)$.
        CENTER_SIZE: Center + size: $(a_c, b_c, c_c, s_a, s_b, s_c)$.
    """

    CORNERS = "corners"
    CENTER_SIZE = "center_size"