SomeOf

class SomeOf(Transform):
    """Apply a random subset of the given transforms.

    When applied to a batch with `per_instance=True` (the default),
    each batch element independently samples its own subset. This
    requires shape- and schema-preserving transforms so the elements
    can be re-stacked. Pass `per_instance=False` to sample a single
    subset for the whole batch.

    Args:
        transforms: Sequence of candidate transforms.
        num_transforms: How many transforms to apply. An `int` for a
            fixed count, or a `(min, max)` tuple to sample the count
            uniformly from that range.
        replace: If `True`, sample with replacement (the same
            transform may be applied more than once).
        **kwargs: See [`Transform`][torchio.Transform] for additional
            keyword arguments.

    Examples:
        >>> import torchio as tio
        >>> augmentation = tio.SomeOf(
        ...     [tio.Noise(), tio.Flip(), tio.Noise(std=0.5)],
        ...     num_transforms=2,
        ... )
    """

    def __init__(
        self,
        transforms: Sequence[Transform] | None = None,
        *,
        num_transforms: int | tuple[int, int] = 1,
        replace: bool = False,
        **kwargs: Any,
    ) -> None:
        super().__init__(**kwargs)
        self.transforms = list(transforms) if transforms else []
        self.num_transforms = num_transforms
        self.replace = replace

    @property
    def _min_n(self) -> int:
        if isinstance(self.num_transforms, int):
            return self.num_transforms
        return self.num_transforms[0]

    @property
    def _max_n(self) -> int:
        if isinstance(self.num_transforms, int):
            return self.num_transforms
        return self.num_transforms[1]

    def forward(self, data):
        if self.copy:
            data = copy.deepcopy(data)
        batch, unwrap = self._wrap(data)
        # The input is copied once above, so children apply without copying.
        with _disabled_copy(self.transforms):
            if self.per_instance and batch.batch_size > 1:
                return unwrap(self._forward_per_element(batch))
            if torch.rand(1).item() >= self.p:
                return unwrap(batch)
            batch = self._apply_subset(batch)
            return unwrap(batch)

    def _apply_subset(self, batch):
        """Apply a randomly chosen subset of transforms to *batch*."""
        n = int(torch.randint(self._min_n, self._max_n + 1, size=(1,)).item())
        n_transforms = len(self.transforms)
        if self.replace:
            indices = torch.randint(0, n_transforms, (n,))
        else:
            n = min(n, n_transforms)
            indices = torch.randperm(n_transforms)[:n]
        for idx in indices:
            batch = self.transforms[idx](batch)
        return batch

    def _forward_per_element(self, batch):
        """Apply an independently chosen subset to each batch element."""
        if self.p == 0:
            return batch
        out_subjects = []
        any_applied = False
        for subject in batch.unbatch():
            if torch.rand(1).item() < self.p:
                any_applied = True
                subject = _apply_to_element(subject, self._apply_subset)
            out_subjects.append(subject)
        if not any_applied:
            return batch
        return _rebatch_with_history(out_subjects, "SomeOf")

    def to_hydra(self) -> dict[str, Any]:
        cfg = super().to_hydra()
        cfg["transforms"] = [t.to_hydra() for t in self.transforms]
        return cfg