Spatial

class Spatial(SpatialTransform):
    r"""Apply resampling, affine motion, and elastic deformation together.

    This transform can:

    1. resample to a new space,
    2. apply a global affine mapping, and
    3. apply a dense elastic field,

    using a single sampling grid.

    The convenience wrappers [`Resample`][torchio.Resample],
    [`Affine`][torchio.Affine], and
    [`ElasticDeformation`][torchio.ElasticDeformation] expose subsets
    of these parameters with sensible defaults.

    Args:
        target: Output space.  Can be one of:

            - A scalar or 3-tuple of floats: output voxel spacing in mm.
              E.g., `1` for 1 mm isotropic, `(0.5, 0.5, 2.0)` for
              anisotropic.
            - A random spacing spec, sampled once per call at apply time: a
              2-tuple `(lo, hi)` (uniform, per axis), a 6-tuple
              `(lo1, hi1, lo2, hi2, lo3, hi3)` (per-axis ranges), a `Choice`,
              or a `torch.distributions.Distribution`. E.g.,
              `target=(1, 2)` or `target=(2, 4, 2, 4, 3, 6)`.
            - A `str`: either a path to an image file, or the name of
              an image in the subject (e.g., `"t1"`).
            - An [`Image`][torchio.Image] instance.
            - A `(spatial_shape, affine)` pair.
            - `None` (default): the output grid matches the input grid.
        scales: Scale factors $(s_1, s_2, s_3)$ for each axis.
            If a single value $x$ is given, all axes use $x$.
            If two values $(a, b)$ are given,
            $s_i \sim \mathcal{U}(a, b)$.
            If six values $(a_1, b_1, a_2, b_2, a_3, b_3)$ are given,
            $s_i \sim \mathcal{U}(a_i, b_i)$ independently.
            A `torch.distributions.Distribution` may also be passed.
            For example, `scales=0.5` halves the apparent object
            size (zoom out), and `scales=2` doubles it (zoom in).
        degrees: Euler rotation angles $(\theta_1, \theta_2, \theta_3)$
            in degrees, following the same value/range/distribution
            convention as *scales*.
        translation: Translation $(t_1, t_2, t_3)$ in mm, following
            the same convention.  The direction depends on the image
            orientation: in RAS+, `translation=(10, 0, 0)` shifts
            10 mm to the right.
        isotropic: If `True`, sample a single scale factor and
            reuse it for all three axes.  *scales* must then be a
            scalar or 2-value range.
        center: Pivot point for rotation and scaling.
            `"image"` (default) uses the image center;
            `"origin"` uses the world-coordinate origin.
        control_points: Optional pre-computed coarse displacement
            field with shape `(n_i, n_j, n_k, 3)` in mm.  If given,
            *num_control_points*, *max_displacement*, and
            *locked_borders* are ignored.
        num_control_points: Number of control points along each
            dimension of the coarse grid.  Can be a single `int`
            (isotropic) or a 3-tuple.  Minimum is 4.  Smaller values
            produce smoother deformations.
        max_displacement: Maximum displacement at each control point,
            in mm.  Follows the same value/range/distribution
            convention as *scales*.  Zero (default) disables elastic
            deformation.
        locked_borders: Number of outer control-point layers whose
            displacement is forced to zero.  `0` keeps all
            displacements; `1` zeros the outermost layer; `2`
            (default) zeros the two outermost layers.
        affine_first: If `True` (default), apply the affine mapping
            before the elastic field.  If `False`, apply the elastic
            field first.  The difference is significant for large
            transforms.
        image_interpolation: `"linear"` (default) or `"nearest"`.
            Used for [`ScalarImage`][torchio.ScalarImage] instances.
        label_interpolation: `"nearest"` (default) or `"linear"`.
            Used for [`LabelMap`][torchio.LabelMap] instances.
        antialias: If `True`, apply Gaussian smoothing before
            downsampling intensity images.  Label maps are never
            smoothed.  The standard deviations follow
            [Cardoso et al., MICCAI 2015](https://link.springer.com/chapter/10.1007/978-3-319-24571-3_81).
        default_pad_value: Fill rule for out-of-bounds intensity
            voxels.  `"minimum"` (default), `"mean"`, `"otsu"`,
            or a numeric value.
        default_pad_label: Numeric fill value for out-of-bounds label
            voxels.
        **kwargs: See [`Transform`][torchio.Transform].

    Note:
        The randomizable parameters (`scales`, `degrees`, `translation`,
        `max_displacement`, and the spacing form of `target`) follow a common
        value/range/distribution convention: a scalar is deterministic, a
        2-tuple $(a, b)$ samples uniformly, a 3-tuple sets per-axis values, a
        6-tuple sets per-axis ranges, and a `Choice` or
        `torch.distributions.Distribution` samples from a discrete set or a
        distribution. Structural parameters (`center`, interpolation, padding)
        are not randomizable.

    Examples:
        >>> import torchio as tio
        >>> # Resample to 1 mm isotropic with a random rotation
        >>> transform = tio.Spatial(
        ...     target=1,
        ...     degrees=(-10, 10),
        ...     translation=(-5, 5),
        ... )
        >>> # Elastic deformation only
        >>> transform = tio.Spatial(
        ...     max_displacement=7.5,
        ...     num_control_points=7,
        ... )
        >>> transformed = transform(subject)
    """

    def __init__(
        self,
        *,
        target: TypeTarget = None,
        scales: TypeParameterValue = 1.0,
        degrees: TypeParameterValue = 0.0,
        translation: TypeParameterValue = 0.0,
        isotropic: bool = False,
        center: TypeCenter = "image",
        control_points: TypeControlPoints | None = None,
        num_control_points: int | TypeThreeInts = 7,
        max_displacement: TypeParameterValue = 0.0,
        locked_borders: int = 2,
        affine_first: bool = True,
        image_interpolation: TypeInterpolation = "linear",
        label_interpolation: TypeInterpolation = "nearest",
        antialias: bool = False,
        default_pad_value: TypePadValue | float = "minimum",
        default_pad_label: int | float = 0,
        **kwargs: Any,
    ) -> None:
        super().__init__(**kwargs)
        self.target = target
        _validate_isotropic(scales, isotropic)
        self.scales = _to_positive_range(scales)
        self.degrees = _to_parameter_range(degrees)
        self.translation = _to_parameter_range(translation)
        self.isotropic = isotropic
        self.center = _parse_center(center)
        self.control_points = (
            _parse_control_points(control_points)
            if control_points is not None
            else None
        )
        self.num_control_points = _parse_num_control_points(num_control_points)
        self.max_displacement = _to_nonnegative_parameter_range(max_displacement)
        self.locked_borders = _parse_locked_borders(locked_borders)
        if self.locked_borders == 2 and 4 in self.num_control_points:
            msg = (
                "locked_borders=2 with 4 control points along any axis yields an"
                " identity elastic field"
            )
            raise ValueError(msg)
        self.affine_first = affine_first
        self.image_interpolation = _parse_interpolation(image_interpolation)
        self.label_interpolation = _parse_interpolation(label_interpolation)
        self.antialias = antialias
        self.default_pad_value = _parse_default_pad_value(default_pad_value)
        if not isinstance(default_pad_label, Number):
            msg = f"default_pad_label must be numeric, got {type(default_pad_label)}"
            raise TypeError(msg)
        self.default_pad_label = float(default_pad_label)

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

    @property
    def supports_per_instance_p(self) -> bool:
        # Per-element gating requires a shape-preserving transform so the
        # gated-out element is a true no-op. A resampling target changes
        # the grid, so keep batch-wide gating in that case.
        return self.target is None

    def _sample_one(
        self,
        shape: TypeThreeInts,
        affine: AffineMatrix,
    ) -> tuple[
        np.ndarray | None,
        Tensor | None,
        tuple[float, float, float] | None,
        bool,
    ]:
        """Sample one set of affine and elastic parameters.

        Args:
            shape: Spatial shape used to build the affine and to suppress
                out-of-plane components for 2D inputs.
            affine: Affine of the reference image.

        Returns:
            A tuple `(forward_affine, control_points, max_displacement,
            has_geometry)` where `has_geometry` is `True` when either an
            affine or an elastic component is present.
        """
        sampled_scales = _sample_scales(self.scales, self.isotropic)
        sampled_degrees = self.degrees.sample()
        sampled_translation = self.translation.sample()
        has_affine = _has_affine_component(
            sampled_scales,
            sampled_degrees,
            sampled_translation,
        )
        control_points, max_displacement = _resolve_control_points(
            self.control_points,
            self.num_control_points,
            self.max_displacement,
            self.locked_borders,
        )
        has_elastic = control_points is not None
        forward_affine = None
        if has_affine:
            forward_affine = _build_forward_affine(
                scales=sampled_scales,
                degrees=sampled_degrees,
                translation=sampled_translation,
                center=self.center,
                shape=shape,
                affine=affine,
            )
        return (
            forward_affine,
            control_points,
            max_displacement,
            (has_affine or has_elastic),
        )

    def make_params(self, batch: SubjectsBatch) -> dict[str, Any]:
        """Sample random parameters and resolve the output space.

        Scales, degrees, translation, and control-point displacements are
        sampled per batch element when per-instance augmentation is
        active (the default for batches), and once otherwise.

        Returns:
            Dict of serializable parameters for `apply_transform` and
            history replay.
        """
        images = self._get_images(batch)
        if not images:
            return {"selected_images": []}

        _, first_batch = next(iter(images.items()))
        first_shape = _get_spatial_shape(first_batch)
        first_affine = first_batch.affines[0]

        params: dict[str, Any] = {
            "selected_images": list(images.keys()),
            "original": _serialize_space((first_shape, first_affine)),
            "affine_first": self.affine_first,
            "image_interpolation": self.image_interpolation,
            "label_interpolation": self.label_interpolation,
            "antialias": self.antialias,
            "default_pad_value": self.default_pad_value,
            "default_pad_label": self.default_pad_label,
        }

        n = self._resolve_n(batch)
        if n is None:
            forward_affine, control_points, max_displacement, has_geometry = (
                self._sample_one(first_shape, first_affine)
            )
            if has_geometry:
                _check_shared_space(images, first_shape, first_affine)
            # Resolve the (possibly random) target after sampling the
            # geometry so the RNG stream matches the batch-shared path.
            target_space = _resolve_target_space(
                self.target,
                batch,
                first_shape,
                first_affine,
            )
            params["target"] = _serialize_space(target_space)
            params["affine_matrix"] = _serialize_matrix(forward_affine)
            params["control_points"] = _serialize_control_points(control_points)
            params["max_displacement"] = (
                list(max_displacement) if max_displacement else None
            )
            return params

        keep = self._keep_mask(batch, n)
        affine_list, control_points_list, displacement_list, any_geometry = (
            self._sample_per_element_geometry(n, first_shape, first_affine, keep)
        )
        if any_geometry:
            _check_shared_space(images, first_shape, first_affine)
        target_space = _resolve_target_space(
            self.target,
            batch,
            first_shape,
            first_affine,
        )
        params["target"] = _serialize_space(target_space)
        params["affine_matrix"] = affine_list
        params["control_points"] = control_points_list
        params["max_displacement"] = displacement_list
        self._tag_batched(
            params,
            batch,
            n,
            keep,
            ["affine_matrix", "control_points", "max_displacement"],
        )
        return params

    def _sample_per_element_geometry(
        self,
        n: int,
        first_shape: TypeThreeInts,
        first_affine: AffineMatrix,
        keep: Tensor | None,
    ) -> tuple[list[Any], list[Any], list[Any], bool]:
        """Sample serializable geometry for each batch element.

        Gated-out elements (``keep[index]`` is false) contribute ``None``
        placeholders so the per-element lists stay aligned with the batch.

        Args:
            n: Number of batch elements.
            first_shape: Spatial shape of the first image.
            first_affine: Affine of the first image.
            keep: Per-element keep mask, or ``None`` to keep all elements.

        Returns:
            The affine, control-point and max-displacement lists plus a
            flag indicating whether any element produced a geometry
            transform.
        """
        affine_list: list[Any] = []
        control_points_list: list[Any] = []
        displacement_list: list[Any] = []
        any_geometry = False
        for index in range(n):
            kept = keep is None or bool(keep[index])
            if not kept:
                affine_list.append(None)
                control_points_list.append(None)
                displacement_list.append(None)
                continue
            forward_affine, control_points, max_displacement, has_geometry = (
                self._sample_one(first_shape, first_affine)
            )
            any_geometry = any_geometry or has_geometry
            affine_list.append(_serialize_matrix(forward_affine))
            control_points_list.append(_serialize_control_points(control_points))
            displacement_list.append(
                list(max_displacement) if max_displacement else None
            )
        return affine_list, control_points_list, displacement_list, any_geometry

    def apply_transform(
        self,
        batch: SubjectsBatch,
        params: dict[str, Any],
    ) -> SubjectsBatch:
        """Apply the spatial mapping to every selected image in *batch*.

        One sampling grid is built per batch element when per-instance
        parameters are present, and a single shared grid otherwise.
        """
        selected_images = params.get("selected_images", [])
        if not selected_images:
            return batch

        target_space = _deserialize_space(params["target"])

        affine_matrix, control_points, max_displacement, per_sample = (
            _resolve_spatial_params(params)
        )
        is_noop = (
            target_space is None
            and affine_matrix is None
            and control_points is None
            and max_displacement is None
            and per_sample is None
        )
        if is_noop:
            # A true no-op (no resampling and no geometry, e.g. every
            # element gated out) must leave the data and the per-sample
            # affines untouched instead of rebuilding an identity grid.
            return batch
        _apply_spatial_to_batch(
            batch=batch,
            image_names=selected_images,
            target_space=target_space,
            affine_matrix=affine_matrix,
            control_points=control_points,
            max_displacement=max_displacement,
            affine_first=params["affine_first"],
            image_interpolation=params["image_interpolation"],
            label_interpolation=params["label_interpolation"],
            antialias=params.get("antialias", False),
            default_pad_value=params["default_pad_value"],
            default_pad_label=float(params["default_pad_label"]),
            per_sample=per_sample,
        )
        return batch

    @property
    def invertible(self) -> bool:
        """Whether this transform can be inverted."""
        return True

    def inverse(self, params: dict[str, Any]) -> _SpatialInverse:
        """Build the inverse transform from recorded parameters.

        The affine component is inverted exactly.  The elastic component
        is approximated by negating the sampled displacement field.  The
        `affine_first` flag is flipped so that the inverse operations
        run in the opposite order. Per-instance parameters are inverted
        element by element.

        Args:
            params: The parameter dict produced by `make_params`.

        Returns:
            A `_SpatialInverse` that resamples back to the original grid.
        """
        original_space = _deserialize_space(params["original"])
        if original_space is None:
            msg = "Spatial inverse needs the original output space"
            raise RuntimeError(msg)

        common: dict[str, Any] = {
            "target": original_space,
            "affine_first": not params["affine_first"],
            "image_interpolation": params["image_interpolation"],
            "label_interpolation": params["label_interpolation"],
            "default_pad_value": params["default_pad_value"],
            "default_pad_label": float(params["default_pad_label"]),
            "copy": False,
            # Invert only the images the forward pass actually transformed,
            # so excluded images (e.g. label maps) are not resampled.
            "include": params["selected_images"],
        }

        batched_keys = params.get("_batched_keys") or []
        if "affine_matrix" in batched_keys:
            per_sample = _invert_per_sample(params)
            return _SpatialInverse(
                affine_matrix=None,
                control_points=None,
                per_sample=per_sample,
                **common,
            )

        affine_matrix = _deserialize_matrix(params["affine_matrix"])
        inverse_affine = None
        if affine_matrix is not None:
            inverse_affine = np.linalg.inv(affine_matrix)
        control_points = _deserialize_control_points(params["control_points"])
        inverse_control_points = None
        if control_points is not None:
            inverse_control_points = -control_points
        return _SpatialInverse(
            affine_matrix=inverse_affine,
            control_points=inverse_control_points,
            **common,
        )