Flip

class Flip(SpatialTransform):
    r"""Reverse the order of elements in an image along the given axes.

    Args:
        axes: Index or tuple of indices of the spatial dimensions along
            which the image might be flipped. Integers must be in
            `{0, 1, 2}`. Anatomical labels may also be used, such as
            `'Left'`, `'Right'`, `'Anterior'`, `'Posterior'`,
            `'Inferior'`, `'Superior'`. Only the first letter of
            the string is used. Anatomical labels are resolved using
            the image orientation.
        flip_probability: Probability that each axis will be flipped
            (per-axis coin flip). This is independent of the `p`
            parameter, which gates the entire transform.
        **kwargs: See [`Transform`][torchio.Transform] for additional
            keyword arguments.

    Tip:
        Specifying the axes as anatomical labels is useful when the
        image orientation is not known.

    Examples:
        >>> import torchio as tio
        >>> # Flip along the first spatial axis
        >>> transform = tio.Flip(axes=0)
        >>> # Flip along the lateral axis (anatomical label)
        >>> transform = tio.Flip(axes='LR')
        >>> # Random per-axis flip with 50% chance each
        >>> transform = tio.Flip(axes=(0, 1, 2), flip_probability=0.5)
    """

    def __init__(
        self,
        *,
        axes: int | str | Sequence[int | str] = 0,
        flip_probability: float = 1.0,
        **kwargs: Any,
    ) -> None:
        super().__init__(**kwargs)
        self.axes = axes
        if not 0 <= flip_probability <= 1:
            msg = f"flip_probability must be in [0, 1], got {flip_probability}"
            raise ValueError(msg)
        self.flip_probability = flip_probability

    def make_params(
        self,
        batch: SubjectsBatch,
    ) -> dict[str, Any]:
        images = self._get_images(batch)
        if not images:
            return {"axes": ()}
        first_img = next(iter(images.values()))

        n = self._resolve_n(batch)
        if n is None:
            orientation = None
            if first_img.batch_size > 0:
                orientation = first_img[0].orientation
            resolved = _resolve_axes(self.axes, orientation)
            flip_mask = torch.rand(3) < self.flip_probability
            axes_to_flip = tuple(a for a in resolved if flip_mask[a].item())
            return {"axes": axes_to_flip}

        keep = self._keep_mask(batch, n)
        axes_list = self._sample_per_element_axes(n, first_img, keep)
        params = {"axes": axes_list}
        self._tag_batched(params, batch, n, keep, ["axes"])
        return params

    def _sample_per_element_axes(
        self,
        n: int,
        first_img: ImagesBatch,
        keep: torch.Tensor | None,
    ) -> list[list[int]]:
        """Sample the flip axes for each batch element.

        Gated-out elements (``keep[index]`` is false) get an empty axis
        list so they are left unflipped.

        Args:
            n: Number of batch elements.
            first_img: First selected image batch, used for per-element
                orientation.
            keep: Per-element keep mask, or ``None`` to keep all elements.

        Returns:
            One list of spatial axes (in ``{0, 1, 2}``) per element.
        """
        axes_list: list[list[int]] = []
        for index in range(n):
            if keep is not None and not keep[index]:
                axes_list.append([])
                continue
            # Resolve anatomical axes per element: each sample may have its
            # own orientation in a batch with per-sample affines.
            resolved = _resolve_axes(self.axes, first_img[index].orientation)
            flip_mask = torch.rand(3) < self.flip_probability
            axes_list.append([a for a in resolved if flip_mask[a].item()])
        return axes_list

    @property
    def supports_per_instance_params(self) -> bool:
        return True

    @property
    def supports_per_instance_p(self) -> bool:
        return True

    def apply_transform(
        self,
        batch: SubjectsBatch,
        params: dict[str, Any],
    ) -> SubjectsBatch:
        axes = params["axes"]
        if self._is_per_instance_params(params):
            for _name, img_batch in self._get_images(batch).items():
                img_batch.data = _flip_per_element(img_batch.data, axes)
            return batch
        if not axes:
            return batch
        dims = [a - 3 for a in axes]
        for _name, img_batch in self._get_images(batch).items():
            img_batch.data = torch.flip(img_batch.data, dims)
        return batch

    @property
    def invertible(self) -> bool:
        return True

    def inverse(self, params: dict[str, Any]) -> Flip | _FlipInverse:
        """Flip is its own inverse."""
        if self._is_per_instance_params(params):
            return _FlipInverse(axes_per_element=params["axes"], copy=False)
        return Flip(axes=params["axes"], copy=False)