Transforms

class Transform(nn.Module):
    """Abstract class for all TorchIO transforms.

    When called, the input can be an instance of
    [`Subject`][torchio.Subject],
    [`Image`][torchio.Image],
    [`torch.Tensor`][torch.Tensor],
    [`numpy.ndarray`][numpy.ndarray],
    [`SimpleITK.Image`](https://simpleitk.org/doxygen/latest/html/classitk_1_1simple_1_1Image.html),
    [`nibabel.Nifti1Image`](https://nipy.org/nibabel/reference/nibabel.nifti1.html),
    [`dict`][dict] containing 4D tensors as values,
    [`ImagesBatch`][torchio.ImagesBatch], or
    [`SubjectsBatch`][torchio.SubjectsBatch].
    The output type always matches the input type.

    All subclasses must override
    [`apply_transform()`][torchio.Transform.apply_transform],
    which receives a [`SubjectsBatch`][torchio.SubjectsBatch] and
    returns the transformed batch.

    Args:
        p: Probability that this transform will be applied. When
            per-instance probability is active (see `per_instance`),
            this is instead the per-element probability and each batch
            element is gated independently.
        copy: Make a deep copy of the input before applying the
            transform. When transforms are composed with
            [`Compose`][torchio.Compose], the outer `Compose`
            copies once and sets `copy=False` on inner transforms
            to avoid redundant copies.
        per_instance: If `True` (default), transforms that support it
            sample independent parameters for each element of a batch
            (and gate each element independently with `p`). If
            `False`, a single parameter set is sampled and applied
            identically to every element, reproducing the legacy
            batch-shared behavior. Single-element inputs (including a
            single [`Subject`][torchio.Subject]) are unaffected by this
            flag.
        include: Sequence of strings with the names of the only images
            to which the transform will be applied.
        exclude: Sequence of strings with the names of the images to
            which the transform will *not* be applied.
    """

    def __init__(
        self,
        *,
        p: float = 1.0,
        copy: bool = True,
        per_instance: bool = True,
        include: list[str] | None = None,
        exclude: list[str] | None = None,
    ) -> None:
        super().__init__()
        if not 0 <= p <= 1:
            msg = f"Probability must be in [0, 1], got {p}"
            raise ValueError(msg)
        self.p = p
        self.copy = copy
        self.per_instance = per_instance
        self.include = include
        self.exclude = exclude

    def __init_subclass__(cls, **kwargs: Any) -> None:
        super().__init_subclass__(**kwargs)
        _TRANSFORM_REGISTRY[cls.__name__] = cls

    def _warn_if_noop(self, *, is_noop: bool, hint: str) -> None:
        """Warn that the transform leaves the data unchanged.

        Augmentation transforms whose parameters are sampled from a
        range default to an identity (no-op) when constructed with no
        arguments, so that randomness must be requested explicitly. This
        warns the user when that happens (or whenever the given
        parameters produce a no-op).

        Args:
            is_noop: Whether the configured transform is an identity.
            hint: Example argument to suggest in the warning message.
        """
        if is_noop:
            warnings.warn(
                f"{type(self).__name__} is a no-op with the given parameters"
                " and will not change the data. Pass arguments to apply an"
                f" effect (e.g. {hint}), or a range like (a, b) for random"
                " augmentation.",
                stacklevel=3,
            )

    def __repr__(self) -> str:
        """Show only non-default fields for a compact repr."""
        from .parameter_range import _ParameterRange

        parts = []
        for name, default in _collect_init_params(type(self)).items():
            value = getattr(self, name, default)
            if isinstance(value, _ParameterRange):
                if value._original == default:
                    continue
            elif value == default:
                continue
            parts.append(f"{name}={value!r}")
        return f"{type(self).__name__}({', '.join(parts)})"

    def __add__(self, other: object) -> Transform:
        """Compose two transforms: `t1 + t2` → `Compose([t1, t2])`."""
        if not isinstance(other, Transform):
            return NotImplemented
        from .compose import Compose

        left = self.transforms if isinstance(self, Compose) else [self]
        right = other.transforms if isinstance(other, Compose) else [other]
        return Compose([*left, *right])

    def __or__(self, other: object) -> Transform:
        """Random choice: `t1 | t2` → `OneOf([t1, t2])`."""
        if not isinstance(other, Transform):
            return NotImplemented
        from .compose import OneOf

        left = self.transforms if isinstance(self, OneOf) else [self]
        right = other.transforms if isinstance(other, OneOf) else [other]
        return OneOf([*left, *right])

    @overload
    def forward(self, data: Subject) -> Subject: ...
    @overload
    def forward(self, data: Image) -> Image: ...
    @overload
    def forward(self, data: Tensor) -> Tensor: ...
    @overload
    def forward(self, data: np.ndarray) -> np.ndarray: ...
    @overload
    def forward(self, data: sitk.Image) -> sitk.Image: ...
    @overload
    def forward(self, data: nib.Nifti1Image) -> nib.Nifti1Image: ...
    @overload
    def forward(self, data: dict) -> dict: ...
    @overload
    def forward(self, data: ImagesBatch) -> ImagesBatch: ...
    @overload
    def forward(self, data: SubjectsBatch) -> SubjectsBatch: ...

    def forward(self, data: Any) -> Any:
        """Apply the transform.

        The output type always matches the input type.

        Args:
            data: Input data to transform.
        """
        if self.copy:
            data = _copy.deepcopy(data)
        batch, unwrap = self._wrap(data)
        # When per-element gating is active, the transform handles the
        # probability itself (masked-out elements get identity params),
        # so skip the batch-wide coin flip here. Apply iff rand < p, so
        # p=0 is always a no-op and p=1 always applies.
        if not self._per_instance_p_active(batch) and torch.rand(1).item() >= self.p:
            return unwrap(batch)
        params = self.make_params(batch)
        batch = self.apply_transform(batch, params)
        # Record history on the batch, unless every element was gated out by
        # per-element probability: that is an exact no-op, and recording it
        # would let history replay (e.g. an invertible spatial transform)
        # trigger an unnecessary identity resample.
        if not _all_elements_gated_out(params):
            trace = AppliedTransform(name=type(self).__name__, params=params)
            if not hasattr(batch, "applied_transforms"):
                batch.applied_transforms = []
            batch.applied_transforms.append(trace)
        result = unwrap(batch)
        # Propagate history to outputs that can carry it
        if (
            hasattr(batch, "applied_transforms")
            and not isinstance(result, (SubjectsBatch, Tensor, np.ndarray))
            and not isinstance(result, dict)
        ):
            with contextlib.suppress(AttributeError):
                result.applied_transforms = list(batch.applied_transforms)
        return result

    @property
    def supports_per_instance_params(self) -> bool:
        """Whether this transform can sample parameters per batch element.

        Defaults to `False`. Transforms that implement per-instance
        parameter sampling override this to return `True`. When `False`,
        the transform always uses batch-shared parameters regardless of
        the `per_instance` flag, preserving the legacy behavior.
        """
        return False

    @property
    def supports_per_instance_p(self) -> bool:
        """Whether this transform can gate each batch element independently.

        Defaults to `False`. Shape-preserving transforms that implement
        per-element probability override this to return `True`.
        Shape-changing transforms must leave it `False` because masked
        and unmasked elements would have incompatible shapes.
        """
        return False

    def _per_instance_active(self, batch: SubjectsBatch) -> bool:
        """Whether per-instance parameter sampling applies to *batch*.

        Per-instance sampling only kicks in for genuine batches
        (`batch_size > 1`); single-element inputs always use the legacy
        scalar path.
        """
        return (
            self.per_instance
            and self.supports_per_instance_params
            and batch.batch_size > 1
        )

    def _per_instance_p_active(self, batch: SubjectsBatch) -> bool:
        """Whether per-element probability gating applies to *batch*."""
        return (
            self.per_instance
            and self.supports_per_instance_p
            and batch.batch_size > 1
            and 0.0 < self.p < 1.0
        )

    def _resolve_n(self, batch: SubjectsBatch) -> int | None:
        """Return the number of parameter sets to sample.

        Returns:
            The batch size when per-instance sampling is active,
            otherwise `None` (the legacy single-sample path).
        """
        return batch.batch_size if self._per_instance_active(batch) else None

    def _keep_mask(
        self,
        batch: SubjectsBatch,
        n: int | None,
    ) -> Tensor | None:
        """Sample a per-element keep mask for per-instance probability.

        Args:
            batch: The batch being transformed.
            n: The resolved number of parameter sets (from
                `_resolve_n`).

        Returns:
            A boolean tensor of shape `(n,)` where `True` marks
            elements that receive the transform, or `None` when
            per-element gating is not active (all elements are kept).
        """
        if n is None or not self._per_instance_p_active(batch):
            return None
        return torch.rand(n) < self.p

    @staticmethod
    def _mask_identity(
        value: Tensor | float,
        keep: Tensor | None,
        *,
        identity: float,
    ) -> Tensor | float:
        """Replace masked-out elements of *value* with an identity value.

        Args:
            value: Sampled parameter, either a scalar (legacy path) or a
                `(B,)` tensor (per-instance path).
            keep: Per-element keep mask, or `None` for no masking.
            identity: The value that makes the transform a no-op for an
                element (for example `0.0` for additive or log-space
                parameters).

        Returns:
            The masked parameter.
        """
        if keep is None or not isinstance(value, Tensor):
            return value
        return torch.where(keep, value, torch.full_like(value, identity))

    @staticmethod
    def _serialize_param(value: Tensor | Any) -> Any:
        """Convert a possibly-tensor parameter to a JSON-serializable form."""
        if isinstance(value, Tensor):
            return value.tolist()
        return value

    @staticmethod
    def _is_per_instance_params(params: dict[str, Any]) -> bool:
        """Whether *params* holds per-element (batched) values."""
        return "_batched_keys" in params

    def _tag_batched(
        self,
        params: dict[str, Any],
        batch: SubjectsBatch,
        n: int | None,
        keep: Tensor | None,
        batched_keys: list[str],
    ) -> None:
        """Annotate *params* with per-instance bookkeeping for history.

        Adds the batch size, the names of the per-element keys, and the
        keep mask so that [`SubjectsBatch.unbatch`][torchio.SubjectsBatch.unbatch]
        can split the history per subject.

        Args:
            params: The parameter dict to annotate in place.
            batch: The batch being transformed.
            n: The resolved number of parameter sets.
            keep: The per-element keep mask, or `None`.
            batched_keys: Names of the params that hold one value per
                element.
        """
        if n is None:
            return
        params["_batch_size"] = batch.batch_size
        params["_batched_keys"] = list(batched_keys)
        if keep is not None:
            params["_keep"] = keep.tolist()

    def make_params(self, batch: SubjectsBatch) -> dict[str, Any]:
        """Sample random parameters for this transform.

        Override in subclasses that have random behavior.

        Args:
            batch: A `SubjectsBatch`.

        Returns:
            Dict of sampled parameters.
        """
        return {}

    def apply_transform(
        self,
        batch: SubjectsBatch,
        params: dict[str, Any],
    ) -> SubjectsBatch:
        """Apply the transform with the given parameters.

        Must be overridden by subclasses. Receives a `SubjectsBatch`
        whose `ImagesBatch` entries contain 5D tensors
        `(B, C, I, J, K)`. Use negative indexing (`-3`, `-2`,
        `-1`) for spatial dims.

        Args:
            batch: A `SubjectsBatch` to transform.
            params: Parameters from `make_params`.

        Returns:
            Transformed `SubjectsBatch`.
        """
        raise NotImplementedError

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

    def inverse(self, params: dict[str, Any]) -> Transform:
        """Return a transform that undoes this one.

        Override in invertible subclasses. The returned transform,
        when applied, reverses the effect of the forward pass with
        the given parameters.

        Args:
            params: The parameters recorded in the forward pass.

        Returns:
            A new `Transform` instance that inverts this one.
        """
        msg = f"{type(self).__name__} is not invertible"
        raise NotImplementedError(msg)

    def _get_images(self, batch: SubjectsBatch) -> dict[str, ImagesBatch]:
        """Get image batches filtered by include/exclude."""
        images = batch.images
        if self.include is not None:
            images = {k: v for k, v in images.items() if k in self.include}
        if self.exclude is not None:
            images = {k: v for k, v in images.items() if k not in self.exclude}
        return images

    def to_hydra(self) -> dict[str, Any]:
        """Export as a Hydra-compatible config dict.

        Returns a dict with `_target_` set to the fully qualified
        class name and only non-default field values included.

        Returns:
            Dict suitable for `hydra.utils.instantiate()`.
        """
        from .parameter_range import _ParameterRange

        cls = type(self)
        target = f"torchio.{cls.__qualname__}"
        cfg: dict[str, Any] = {"_target_": target}

        for name, default in _collect_init_params(cls).items():
            value = getattr(self, name, default)
            if isinstance(value, _ParameterRange):
                if value._original == default:
                    continue
                value = _hydra_value(value._original)
            elif value == default:
                continue
            else:
                value = _hydra_value(value)
            cfg[name] = value
        return cfg

    @staticmethod
    def _wrap(
        data: Any,
    ) -> tuple[Any, Any]:
        """Wrap any input into a SubjectsBatch; return (batch, unwrap_fn)."""
        from ..data.batch import ImagesBatch
        from ..data.batch import SubjectsBatch

        match data:
            case SubjectsBatch():
                return data, _unwrap_subjects_batch
            case ImagesBatch():
                sb = SubjectsBatch({"tio_default_image": data})
                return sb, _unwrap_images_batch
            case Subject():
                sb = SubjectsBatch.from_subjects([data])
                return sb, _unwrap_subject
            case dict():
                return _wrap_dict(data)
            case _:
                return _wrap_scalar_input(data)

class SpatialTransform(Transform):
    """Base for transforms that modify spatial geometry.

    Spatial transforms apply to all images (ScalarImage and LabelMap),
    and also transform any Points and BoundingBoxes attached to the
    Subject.
    """

class IntensityTransform(Transform):
    """Base for transforms that modify voxel intensities.

    Intensity transforms apply only to `ScalarImage` instances,
    leaving `LabelMap` and annotations unchanged.
    """

    def _get_images(self, batch: SubjectsBatch) -> dict[str, ImagesBatch]:
        """Filter to ScalarImage batches only, then apply include/exclude."""
        images = {
            k: v for k, v in batch.images.items() if v._image_class is ScalarImage
        }
        if self.include is not None:
            images = {k: v for k, v in images.items() if k in self.include}
        if self.exclude is not None:
            images = {k: v for k, v in images.items() if k not in self.exclude}
        return images

@dataclass
class AppliedTransform:
    """Record of a transform application, stored in Subject history.

    Attributes:
        name: Class name of the transform.
        params: Sampled parameters (JSON-serializable).
    """

    name: str
    params: dict[str, Any] = field(default_factory=dict)

class To(Transform):
    """Move all data to a device and/or cast to a dtype.

    Wraps the `to()` method as a transform so it can be used inside
    [`Compose`][torchio.Compose] pipelines.

    Args:
        *to_args: Positional arguments forwarded to
            [`torch.Tensor.to()`](https://pytorch.org/docs/stable/generated/torch.Tensor.to.html).
            Typically a device string (`"cpu"`, `"cuda"`,
            `"mps"`) or a `torch.dtype` (`torch.float16`).
        **to_kwargs: Keyword arguments forwarded to
            `torch.Tensor.to()`.

    Examples:
        >>> import torchio as tio
        >>> transform = tio.To(torch.float16)
        >>> transform = tio.To("cuda")
        >>> pipeline = tio.Compose([
        ...     tio.To("cuda"),
        ...     tio.Noise(std=0.1),
        ... ])
    """

    def __init__(self, *to_args: Any, **to_kwargs: Any) -> None:
        super().__init__()
        self.to_args = to_args
        self.to_kwargs = to_kwargs

    def make_params(self, batch: SubjectsBatch) -> dict[str, Any]:
        return {"to_args": self.to_args, "to_kwargs": self.to_kwargs}

    def apply_transform(
        self,
        batch: SubjectsBatch,
        params: dict[str, Any],
    ) -> SubjectsBatch:
        batch.to(*params["to_args"], **params["to_kwargs"])
        return batch